HFAI Blog RSS

FlashAttention：具有 IO 感知，快速且内存高效的新型注意力算法

Vachel — Thu, 19 Oct 2023 21:00:00 GMT

Transformer 模型的核心是自注意力机制（self attention），其在序列长度上时间和存储的复杂度都在 $O(N^2)$ 级别。随着大语言模型（LLMs）规模的不断扩大，为 LLM 配备更长的上下文背景，在工程实现上面临着非常大的挑战。

来自斯坦福大学计算机系与纽约州立大学布法罗分校的科研团队发表了一种新型的注意力算法，名叫 FlashAttention ，其不仅拥有比 PyTorch 标准注意力快 2~4 倍的运行速度，所需内存还减少了 5~20 倍。后续发布的 FlashAttention2，Flash Decoding 还拥有更夸张的性能加速表现。

幻方基础研究研发的大模型训练工具 HAI-LLM 全系采用 FlashAttention，大幅提高显卡利用率，实现了非常优异的训练表现。本系列文章将为大家深入浅出聊聊 FlashAttention 背后的技术和我们的实践经验。

论文地址：https://arxiv.org/abs/2205.14135

项目源码：https://github.com/Dao-AILab/flash-attention

背景

传统的注意力算法其内存效率是 $O(N^2)$ 的。过去一些优化注意力机制的方法是采用近似值，例如稀疏近似、低秩近似以及它们的组合。尽管这些方法可以将计算降低到线性或接近线性（ $O(N)$ ），但它们过于关注降低每秒所执行的浮点运算次数（FLops），并且倾向于忽略来自内存访问（IO）的开销。

多年来 GPU FLOPS 的增长速度一直比内存吞吐量（TB/s）的增长更快。我们在 A100 或同级别显卡上优化模型训练的实践中发现，内存吞吐量才是影响训练进一步提效的重要瓶颈。FLOPS 和内存吞吐量需要紧密结合，才能充分提高的训练效率。这就需要我们在软件层面上进行更加细致的设计。

如下图所示：

上图展示了 CPU 和 GPU 不同层级内存的吞吐量和容量。可以看到内存不是一个单一的部件，它在本质上是分层的，一般的规则是：内存越快，越昂贵，容量越小。

以 A100 为例：A100 GPU 有 40~80GB 的高带宽内存（HBM），带宽为 1.5-2.0 TB/s，而每 108 个流处理器有 192KB 的 SRAM，带宽估计在 19TB/s 左右。可以看到虽然 SRAM 容量小了很多，但是速度却提升了10倍，所以如何高效的利用 SRAM 是提速注意力算法的关键。

标准注意力算法

我们首先看看标准注意力算法背后的计算逻辑：

可以看到标准注意力算法基本上将 HBM 加载/存储操作视为0成本（它并不能感知 IO）。

下图展示了 GPT-2 模型中一个 Attention 算子的完整计算耗时统计：

可以看到，masking，softmax 和 dropout 操作占用了大量时间，而主要利用 FLOPS 的矩阵乘法（Matmul）却只占用了一部分时间。因此，感知硬件 IO 进行优化的 FlashAttention 算法被提出，其可以大幅减少冗余的 HBM IO 并充分利用 SRAM 进行计算加速。

FlashAttention

FlashAttention 思路是：既然标准注意力算法要将 S 写回 HBM，而这个步骤只为了重新加载计算 Softmax，那么我们可以将其保存在 SRAM 中，等执行完所有中间步骤后，再将最终结果写回 HBM。如下图所示：

可以看到 FlashAttention 将多个操作融合在一起，其只从 HBM 加载一次，执行融合的算子操作，然后将结果写回 HBM。融合操作主要采用了如下两种技术：

Tiling：矩阵分块计算，在不访问整个输入的情况下计算 Softmax 函数的缩减，在前向和后向传播时都使用；
Recomputation：时间换空间，不存储中间注意力矩阵而采用重计算的方式，仅在后向传播时使用。

完整的伪代码如下：

1. Tiling 分块计算

对于有限的 SRAM 容量， $N^2$ 的存储用量使得序列长度（N）限定在了一定范围，因此我们要进行矩阵分块计算。对于矩阵乘法与逐点操作（scale，masking，dropout）的分块计算是比较容易实现的，主要障碍是 Softmax 函数，因为其需要将所有的分数列耦合在一起。为此研究者使用了一个技巧：既然 Softmax 与注意力 $\mathbf{K}$ 的列是耦合的，通过引入了两个额外的统计量 $m(x),l(x)$ 来进行解耦，实现了分块计算。具体如下：

$m(x):=\max_i x_i, ~~f(x):=[e^{x_i-m(x)}...e^{x_B-m(x)}], ~~l(x):=\sum_i f(x)_i, ~~softmax(x):=\frac{f(x)}{l(x)}$

对于两个向量 $x^{(1)},x^{(2)}\in R^B$ ，解耦拼接向量 $x=[x^{(1)},x^{(2)}]\in R^{2B}$ 的 Softmax 计算：

$m(x)=m([x^{(1)},x^{(2)}])=\max(x^{(1)},x^{(2)}), f(x)=[e^{m(x^{(1)})-m(x)}f(x^{(1)}) ~~ e^{m(x^{(2)})-m(x)}f(x^{(2)})]$ $l(x)=l([x^{(1)},x^{(2)}])=e^{m(x^{(1)})-m(x)}l(x^{(1)}) + e^{m(x^{(2)})-m(x)}l(x^{(2)}), ~~ softmax(x)=\frac{f(x)}{l(x)}$

需要注意的是，可以利用 GPU 多线程同时并行计算多个块的 Softmax。为了充分利用硬件性能，多个块的计算不是串行的，而是并行。

2. 重计算

为了避免产生冗余的 HBM 读写次数，FlashAttention 没有为后向传递保存很大的中间结果矩阵。

在标准注意力实现中，后向传递计算 Q,K,V 的梯度时，需要用到 NxN 的中间矩阵 S,P ，但这两个矩阵并没有保存下来。研究用的技巧是重计算，保存了两个统计量 $m(x),l(x)$ ，后向传递时在高速的 SRAM 上快速地重新计算，通过分块的方式重新计算出注意力矩阵 S,P 。这种方式比标准方法要快很多。

实验

相比于标准的注意力算法，FlashAttention 虽然由于反向传播需要重新计算导致 GFLOPs 增加，但是 FlashAttention 有效减少了 HBM 的 I/O ，运行时间显著减少，如下图左所示：

同时从上图右也可以看到，随着 Block Size 增大，HBM 的访问次数减少，运行时间也随之减少。当 Block Size 超过256 时，尽管 HBM 访问次数在减少，但运行时间并没有减少。这时性能受到了其他因素的限制，例如，计算受限。另外需要注意的是，更大的 Block Size 可能会导致执行一次融合算子操作需要的显存超出 SRAM 的大小。

在 A100 显卡上进行实验，FlashAttention 的加速效果如下图所示：

内存的变化为：

可以看到，在不同的序列长度下组合 dropout 和 masking，都有不同程度的加速效果；随着序列长度的增加，FlashAttention 对于内存消耗有着不断优化的效果。

总结

多数大语言模型输入输出的最大序列长度只有 2K 或 4K，本质原因是 transformer 的核心组件 self-attention 块的计算复杂度和空间复杂度是 $O(N^2)$ 的。FlashAttention 的成功启发我们，可以通过分块计算、算子融合和重计算技术，实现深度学习模型的优化与加速，这对于AI工业实践走向深水区有很大的借鉴意义。

StreamingLLM：一种能够接受近乎无限长度文本的大模型框架

Vachel — Thu, 12 Oct 2023 21:00:00 GMT

使 LLM 模型可以处理无限长文本是一项重要挑战，因为存储之前所有的 KV 会受限于内存，模型很难生成超过其训练序列长度的文本。

最近来自 MIT 高性能 AI 计算的研究者们发布了一个能够接受400万 token 的大模型结构，名叫 StreamingLLM，它只保留最新的 token 和 attention sinks，而丢弃中间的 token。这样，模型无需刷新缓存就能从最近的标记生成流畅的文本。这种方式不仅不会降低多少推理速度，而且理论上能够处理无限长度的输入。

StreamingLLM 非常适合用在模型需要持续运行而不需要大量内存或依赖过去数据的场景，例如基于 LLM 的日常助手。截止该稿撰写时，此项目已获得4.6k的 Github Star。本文将和大家一起深入了解这项工作的细节。

论文地址：https://arxiv.org/abs/2309.17453

源码地址：https://github.com/mit-han-lab/streaming-llm

概述

将 LLM 应用于无限输入流时，会遇到两个主要挑战：

缓存冗余：在解码阶段，基于 Transformer 的 LLM 会缓存之前所有 token 的 KV，这会导致内存使用过多，增加解码延迟；
外推困境：现有模型的长度外推能力有限，即当序列长度超过预训练时设定的注意窗口大小时，其性能就会下降。

一种直观的方法被称为 Window Attention，其仅在最近 token 的 KV 状态上维护一个固定大小的滑动窗口。虽然这种方法能保持稳定的内存使用率和解码速度，但一旦序列长度超过缓存大小，开始对 token 的 KV 进行删除，模型的性能就会急剧下降。如下图 b 所示：

图1 StreamingLLM 与现有方法的对比

另一种策略是带重计算的滑动窗口（如上图 c 所示），这种方法会从最近的 token 中重建 KV 状态，虽然它在长文本中表现良好，但其速度明显较慢，在现实应用中并不实用。

研究者探寻 Window Attention 失效的原因，发现了自回归 LLMs 的一个有趣现象：大模型的 Attention 计算中，大量 Attention 权重被分配给某些初始 token，无论这些 token 与语言建模任务是否相关。如下图所示：

图2 LLama2 7B 模型的注意力分数的可视化

研究者们将这些 token 命名为 “attention sinks”。尽管这些初始 token 缺乏语义，但它们收集了显著的注意力分数。研究者将这个现象归因于计算中的 Softmax 操作，它要求注意力分数总和为 1 。因此当前 query 对许多先前 token 没有很强匹配的时候，仍需要在某个与语义不太相关的地方，分配这些不需要的注意力值，来使得最终分数总和为 1。初始 token 作为 sink token 的原因很直观：由于自回归语言建模的特性，初始 token 对几乎所有后续标记都是可见的，这使得它们更容易被训练成 attention sinks。

方法

借助上面的观察，研究者们提出了 StreamingLLM，一种简单而高效的框架。如图1 d 所示，该模型就是把初始 token 的 KV 保留下来，然后再拼上滑动窗口的 KV，就可以保持稳定的模型性能，因为保留 attention sinks 可以让窗口内的注意力分数分布接近正常。

1. Window Attention 失效的原因

在本节中，我们将使用 attention sinks 的概念来解释 Window Attention 失败的原因，并以此详细阐述 StreamingLLM 的灵感来源。

图3 2万 token 下不同 LLM 的语言建模困惑

上图很明显可以得到，当文本长度超过缓存大小时，由于排除了初始标记，困惑度会激增。这表明，无论初始词块与被预测词块的距离如何，初始词块对于保持 LLM 的稳定性都至关重要。

图2 展示了 llama2 7b 模型的所有层和 head 的注意力图。可以发现，除了最下面两层之外，模型在所有层和 head 都始终专注于初始 token。这意味着：在注意力计算中，去掉这些初始 token 的 KV 将去掉 SoftMax 函数（公式 1）中分母的很大一部分。这种改变会导致注意力分数的分布发生显著变化，从而偏离正常推理设置中的预期。

$SoftMax(x)_i = \frac{e^{x_i}}{e^{x_i}+\sum_{j=2}^{N}e^{x_i}}, x_1 \gg x_j, j \in 2,...,N$

那么初始 token 为什么会如此重要，该论文的研究者给出了如下两种可能：

初始 token 的语义至关重要；
或者模型学习偏重于初始 token 的绝对位置

研究者用“\n”替代前四个标记符，发现模型仍然会非常重视初始换行符。这表明，初始token的绝对位置比其语义价值更加重要。

2. 具有 attention sinks 的滑动 KV 缓存

研究者提出的 StreamingLLM 是一种直接的方法，其可以在不进行任何模型微调的情况下恢复窗口注意的困惑度。除了当前的滑动窗口 token 之外，研究者还在注意力计算中重新引入了几个初始 token 的 KV。如下图所示：

图4 StreamingLLM 的 KV 缓存结构

StreamingLLM 中的 KV 缓存在概念上可以分为两部分：

attention sinks（四个initial tokens）稳定注意力计算；
滑动 KV 缓存保留了最新的标记，这对于语言建模至关重要。

在确定相对距离并将位置信息添加到 token 时，StreamingLLM 专注于缓存中的位置，而不是原始文本中的位置。这种区别对于 StreamingLLM 的性能至关重要。例如，如果当前缓存的 token 为 [0, 1, 2, 3, 6, 7, 8]，那么在解码第 9 个令牌的过程中，分配的位置为 [0, 1, 2, 3, 4, 5, 6, 7]，而不是原始文本中的位置 [0, 1, 2, 3, 6, 7, 8, 9]。

3. 带有 attention sinks 的预训练 LLMs

如1小节所述，模型过度关注多个初始 token 的一个重要原因是缺乏一个指定的 sink token 来卸载过多的注意力分数。因此，模型会无意中将全局可见的标记（主要是初始 token）指定为 attention sinks。一种可能的补救办法是有意识地加入一个可训练的全局 attention sinks 标记，称为 “Sink Token”，作为不必要注意力分数的存放处。另外，也可以用类似于 SoftMax-off-by-One 的变体来取代传统的 SoftMax 函数：

$SoftMax_1(x)_i=\frac{e^{x_i}}{1+\sum_{j=1}^{N}e^{x_i}}$

这种 Softmax 方法不要求所有上下文 token 的注意力分数总和为 1。但要注意的是，这种替代方法等同于在注意力计算中使用 KV 特征全为零的 token 。作者将这种方法命名为 “Zero Sink”，以便与其提出的框架保持一致。

实验

研究者使用四个著名的最新模型对 StreamingLLM 进行了评估：Llama-2、MPT、PyThia 和 Falcon。研究者将 StreamingLLM 与密集注意力、窗口注意力和带重计算的滑动窗口方法等既定基准进行比较。

1. 长文本语言建模测试

研究者首先使用 PG19 测试集评估 StreamingLLM 的语言建模复杂度。如之前图3所示，在跨度为 20K token 的文本上，StreamingLLM 的困惑度可以与 Oracle 基准（滑动窗口与重新计算）相媲美。同时，当输入长度超过预训练窗口时，密集注意力技术就会失效，而当输入长度超过缓存大小时，窗口注意力技术就会陷入困境，导致初始标记被剔除。

图5 StreamingLLM 优化的 LLM 在 4M token 文本上语言建模表现

上图进一步证实了 StreamingLLM 可以可靠地处理异常扩展的文本。在 400 多万个 token 上，四种著名的 LLM 依托 StreamingLLM 框架表现出稳定的性能，其困惑度几乎与具有重新计算基线的滑动窗口的困惑度相匹配。

2. Sink Token 的有效性

为了验证方法3中提出的 sink token 的有效性，即在所有预训练样本中引入 sink token 用于改善 StreamingLLM，研究者训练了两种语言模型，每个模型有160万个参数。在相同的条件下，一个模型采用原始训练设置上，另一个在每个训练样本的开头加入了一个 sink token。

图 6 有无 sink token 下的 Loss 下降曲线

如上图所示，与普通模型相比，使用 sink token 训练的模型表现出相似的收敛动态。使用 sink token 预训练的模型与使用 vanilla 方法训练的模型的性能相似。值得注意的是，vanilla 模型需要将多个标记作为 attention sink 相加以保持稳定的困惑度。相比之下，使用 sink token 训练的模型仅使用 sink token 实现了令人满意的性能。

图 7 使用 sink token 后的注意力分数可视

上图展示了使用和不使用 sink token 预训练的模型的注意力图。可以看到，没有 sink token 的模型，类似于 Llama2 7B（如图 2 所示），显示了早期层的局部注意力和更深的层关注初始标记。相比之下，使用 sink token 训练的模型始终集中在层和头部的接收器上，表明有效的注意力卸载机制。

研究者还展示了使用 StreamingLLM 进行流式 QA 的结果，还有一系列消融实验，效率实验等，感兴趣的读者可以阅读论文原文了解其细节。

总结

总的来说，该工作是一个非常典型的观察规律、分析原因、运用到实际方法的工作。其提出的 attention sink，也许像 BERT 中放在句首用于分类的 [CLS] 标记，或者是图神经网络中的 Super Node，接收的是全局或者说整个句子的信息。实际上这种方式不仅仅可以用在 LLM 上，所有 attention-based 的模型都可以参考这一方式将输入序列的长度变得更加灵活。而对于其他 AI 研究领域来说，则是很好的将解释模型与研发模型结合的例子。

Continuous Batching：一种提升 LLM 部署吞吐量的利器

Vachel — Sun, 10 Sep 2023 21:00:00 GMT

由于 LLM 巨大的 GPU 内存开销和计算成本，在大多数应用中，机器学习工程师通常通过内部调整（如量化和对 CUDA 核的定制）来优化。然而，由于 LLM 通过迭代生成其输出，并且 LLM 推理通常涉及内存而不是计算，因此在很多实践中，优化系统级批处理可以使性能差异达到10倍甚至更多。

一种最近提出的优化方法是连续批处理（Continuous batching），也称为动态批处理或基于迭代级的批处理。其具有如下惊人的效果：

基于vLLM，使用连续批处理和连续批处理特定的内存优化，可以实现多达23倍的吞吐量提升；
对于 HuggingFace 本地生成推理，使用连续批处理，可以实现8倍的吞吐量提升；
基于 NVIDIA 的 FasterTransformer，使用优化过的模型实现，可以实现4倍的吞吐量提升。

本博客接下来将详细介绍相关技术的细节。

文章翻译自: https://www.anyscale.com/blog/continuous-batching-llm-inference

LLM 推理

LLM 推理是一个迭代过程，在每个新前馈循环后获得一个额外的完成标记。例如，如果您提示一个句子”What is the capital of California:“，它需要进行十次前馈循环才能得到完整的回答[“S”,“a”,“c”,“r”,“a”,“m”,“e”,“n”,“t”,“o”]。

如下图所示：

上图示例显示了一个支持最大序列长度为8个标记（T1，T2，…，T8）的假设模型。从 Prompt（黄色）开始，迭代过程逐个生成一个标记（蓝色）。一旦模型生成了一个结束序列标记（红色），生成循环停止。这个例子显示了仅包含一个输入序列的 Batch，因此 Batch 大小为1。

现在，我们已经了解了 LLM 推理的迭代过程，接下来我们来探讨一些不熟悉的事情：

LLM 推理是内存 IO 限制，而不是计算限制。换句话说，目前加载 1MB 的数据到 GPU 所需的时间比 1MB 的数据在GPU上计算所需的时间长。这意味着 LLM 推理的吞吐量很大程度上取决于您能将多少批数据装入到高速GPU 内存中；
GPU 内存的消耗量随着基本模型大小和标记长度的增加而增加。如果我们将序列长度限制为 512，那么在一个批处理中，我们最多只能处理28个序列；一个序列长度为 2048 则批处理大小最多只能为7个序列。需要注意的是，这只是一个上限，因为中间计算结果没有留下存储的空间。

这意味着优化内存使用有很多余地。这就是为什么像 AutoGPTQ 这样的模型量化策略具有如此强大的力量的原因；如果您可以将内存使用量减半，您将能够为更大的批处理提供更多的空间。然而，并不是所有的策略都需要对模型权重进行修改。例如，FlashAttention 通过重新组织注意计算来减少内存IO，从而发现显著的性能提升。

连续批处理是另一种内存优化技术，它不需要对模型权重进行修改。接下来，我们解释一下连续批处理如何工作，以及如何提高LLM生成过程的内存效率。

朴素批处理与静态批处理

我们称这种传统的批处理方法为静态批处理，因为批大小在推理完成之前保持不变。以下是LLM推理中静态批处理的示意图：

如上图所示，在第一遍迭代（左）中，每个序列从提示词（黄）中生成一个标记（蓝色）。经过几轮迭代（右）后，完成的序列具有不同的尺寸，因为每个序列在不同的迭代结束时产生不同的结束序列标记（红色）。尽管序列3在两次迭代后完成，但静态批处理意味着 GPU 将在批处理中的最后一个序列完成。

与传统深度学习模型不同，LLM 的批处理方法比较复杂，因为其推理是迭代性的。直观上，这是因为请求可能提前完成，但是释放其资源并添加可能处于不同完成状态的新请求是棘手的。这意味着，随着 GPU 在批次中不同序列的生成长度与最大生成长度不同而未充分利用。在上面的右图，这通过序列1、3和4的序列末标记后的白色方块来说明。

静态批处理不充分利用GPU的频率是多少？这取决于批次中序列的生成长度。例如，可以使用LLM推断来发射单个标记作为分类任务（有更好的方法可以做到这一点，但让我们以此为例）。在这种情况下，每个输出序列具有相同的大小（1个标记）。如果输入序列也具有相同的大小（例如，512个标记），则每个静态批次将实现最佳的GPU利用率。

另一方面，LLM 驱动的聊天机器人服务不能假定固定长度的输入序列，也不能假定固定长度的输出序列。在撰写本文时，专有模型提供的最大上下文长度超过8K个标记。使用静态批处理，生成输出的变化可能会导致GPU严重未充分利用。难怪OpenAI首席执行官Sam Altman称计算成本高得惊人。

如果没有对用户输入和模型输出的限制假设，未经优化的生产级 LLM 系统无法充分利用 GPU 并导致不必要的较高成本。

连续批处理

行业意识到这种批处理的低效性，并提出了更有效的解决方案。OSDI 2022 上发表的 Orca 是第一篇解决这个问题的论文。它采用了迭代级调度，其中批大小根据每次迭代确定。结果是，一旦批中的一个序列完成生成，就可以在其位置插入一个新的序列，从而实现比静态批处理更高的GPU利用率。

如上图所示，使用连续批处理完成七条序列。左图显示了单个迭代后的批，右图显示了多次迭代后的批。一旦一个序列产生结束序列标记，我们在其位置插入新的序列（即序列S5、S6和S7）。这实现了更高的 GPU 利用率，因为 GPU 不需要等待所有序列完成才开始新的一个。

现实情况比这个简化模型更复杂：因为预填充阶段需要计算，并且与生成阶段的计算模式不同，因此它不能很容易地与令牌的生成一起进行批量。连续批处理框架目前通过超参数来管理这个问题：等待已服务比和等待结束序列标记的请求比（waiting_served_ratio）。

目前，HuggingFace 已经将连续批处理作为他们的基于 Rust 和 Python 的文本生成推理服务器的生产版本。我们使用他们的实现来了解连续批处理在 benchmarks 中的性能特征。

实验验证

我们旨在了解连续批量处理与静态批处理在模拟实时推理工作负载上的表现。本质上，我们关心成本。我们将成本分解为吞吐量（Throughput）和延迟（Latency）这两个指标，因为成本直接取决于您在给定延迟下以多高效的方式提供服务的程度。

我们在一个 NVIDIA A100 GPU 上进行吞吐量（Throughput）和延迟（Latency）的基准测试。我们的 A100 有 40GB 的 GPU 内存。我们选择了 Meta的OPT-13B 模型，因为每个测试框架都与此模型配合良好。我们选择了 13B变体，因为它既不需要张量并行，又足够大以产生内存效率挑战。我们选择不使用张量并行，尽管静态批处理和连续批处理都可以使用张量并行。

我们测试了两种静态批处理框架和三种连续批处理框架，他们分别是：

吞吐量测试结果

根据我们对静态批处理的了解，我们猜测在序列长度波动较大的每个批次中，连续批处理的表现会显著更好。为了证明这一点，我们针对每个框架分别运行我们的吞吐量基准测试4次，每次在具有更高序列长度波动的数据集上进行。

为了做到这一点，我们创建了一个包含 1000 个序列，每个序列具有 512 个输入标记的数据集。我们通过忽略端尾标记并设置最大标记配置来始终生成每个请求的序列长度。然后我们生成 1000 个生成长度，每个请求都来自一个由均值为 128 的指数分布产生的样本。我们使用指数分布，因为它是对像 ChatGPT 这样的应用程序中可能遇到的生成长度的一个好近似。为了改变每个运行的方差，我们仅选择小于或等于 32、128、512 和 1536 的样本。总输出序列长度最多为 512+32=544、512+128=640、512+512=1024 和 512+1536=2048。

然后我们使用一个简单的基准测试脚本向我们的模型服务器提交 HTTP 请求。基准测试脚本以突发方式提交所有请求，使计算饱和。

结果如下：

可以看到，静态批处理和朴素连续批处理在较低的序列长度波动下表现大约相同。然而，随着方差增加，静态批处理的性能显著下降，直到生成长度达到 1536 时，性能下降至 81 个字符/秒。FasterTransformers 在静态批处理方面的表现明显优于静态批处理，几乎与连续批处理保持同步，直到生成长度达到 1536。

令人印象深刻的在这里的是vLLM。对于每个数据集，vLLM的性能是朴素连续批处理的两倍以上。我们尚未分析vLLM 性能中优化贡献最大是什么，但我們怀疑 vLLM 动态预留空间的能力，而不是提前预留空间，使得 vLLM 可以显著增加批量大小。

我们将这些性能结果相对静态批处理的关系绘制出来：

需要注意的是，FasterTransformer 的4倍改进虽然令人印象深刻，但我们非常感兴趣在 NVIDIA 实现它时进行基准测试。然而，即使使用优化后的模型，连续批处理在静态批处理方面仍然具有显著的改进。如果您像 vLLM 一样采用连续批处理和迭代级调度来实现进一步的内存优化，性能差距就会变得非常巨大。

延时测试结果

类似地，我们在一个真实的用例上进行延迟-吞吐量权衡优化。我们对每个框架进行基准测试，并测量随着每个框架的累积分布函数（CDF）的变化，该函数如何随着每个框架的部署而变化。

与吞吐量基准测试相同，我们配置模型以始终生成指定数量的每个请求的标记。我们通过从1个标记到512个标记的均匀分布中采样长度来准备100个随机的提示。我们通过从均值为128且最大大小为1536的指数分布中采样100个输出长度来进行预测。这些数字之所以选择，是因为它们在现实世界中是合理的，并允许生成器充分利用我们模型的全部上下文长度（512+1536=2048）。

我们没有像在吞吐量基准测试中那样同时提交所有请求。而是通过预先延迟每个请求一定的时间来进行基准测试。我们通过泊松分布来确定每个请求提交后等待多长时间。泊松分布由λ和预期率（即每秒击中模型端点的查询数量）进行参数化。我们测量在 QPS 为1和 QPS 为4时的延迟，以了解随着负载的变化，延迟分布如何变化。

结果如下：

可以看到，在提高吞吐量的同时，连续批处理系统也改善了中位数延迟。这是因为连续批处理系统允许在有空间的情况下，将新的请求添加到现有的批次中。但是，其他百分位数呢？事实上，我们发现它们在所有百分位数上都改善了延迟：

QPS=1时，每个框架的生成请求中累积分布函数（CDF）的形状都有所不同，这是由于连续批处理系统中的迭代级批调度导致的。所有连续批处理在 QPS=1 时表现约等于静态批处理；与静态批处理相比，FasterTransformers 的表现明显更好。

连续批处理改善所有百分位数的延迟的原因与改善 p50 的延迟原因相同：新的请求可以添加，而不管其他序列在生成批次中多远。然而，与静态批处理一样，连续批处理仍然受到 GPU 可用空间的限制。随着您的服务系统接收到的请求数量饱和，即平均批量大小更高，那么在收到新请求时，立即在 GPU 上注入新请求的机会就越少。我们可以看到，随着平均QPS增加到4，这一点变得更加明显：

QPS=4 时，每个框架的生成请求中累积分布函数（CDF）的形状与 QPS=1 时的情况相似，FasterTransformer 的分布与静态批处理的分布相似。然而，与 QPS=1 相比，FasterTransformer 的曲线更接近于静态批处理，而 text-generation-inference 和 Ray Serve 的连续批处理实现与 FasterTransformer 的曲线相似，但明显不如 vLLM。

我们观察到，随着系统的饱和，FasterTransformer 变得更加接近于静态批处理，而 text-generation-inference 和 Ray Serve 的连续批处理实现正在朝着与 FasterTransformer 的曲线相似的方向发展。换句话说，随着系统的饱和，立即在 GPU 上注入新请求的机会更少，因此请求延迟上升。这与 vLLM 的曲线相吻合：在 QPS=1 和 QPS=4 之间，它的最大批量大约增加了 1500 个token。

总结

LLMs 具有令人惊叹的功能，我们相信其影响仍有很大的发掘空间。我们已经分享了如何使用一种新的服务技术：连续批处理，以及它如何超越静态批处理。它通过减少浪费安排新请求的机会来提高吞吐量，并能够立即将新请求注入计算流中来改善延迟。我们期待人们如何利用连续批处理，以及这个行业将从此走向何方。

LLaMA-2 技术详解（一）：数据打标

Freja/Vachel — Thu, 10 Aug 2023 21:00:00 GMT

LLaMA 是目前备受关注的开源大语言预训练模型。最近 Meta 发布了 LLaMA 2，它是 LLaMA 的下一代版本，具有商业友好的许可证。LLaMA 2 有 3 种不同的尺寸：7B、13B 和 70B，在对话用例下微调的模型版本被称为 LLaMA 2-Chat。Llama-2-Chat 模型在大多数基准测试中都优于现有的开源聊天模型，在对有用性和安全性的人工评估中，LLama-2-Chat 与 ChatGPT，PaLM 等一些流行的封闭源模型不相上下。

LLama-2-Chat 几乎是现有开源模型中唯一做了 RLHF 的模型。LLaMA-2 经过 5 轮 RLHF 后，在 Meta 自己的 Reward 模型与 GPT-4 的评价下，都表现出了超过 ChatGPT 性能。

本文将为大家整理 LLaMA-2 论文中提到用来训练 Reward Model 的数据获取流程。

标注流程

LLaMA-2 的打标主要考虑有用性(Helpfulness)和安全性(Safety)

有用性(Helpfulness)：response 对 prompt 中问题回答得有多好，是否提供了需要的信息
安全性(Safety)：回答是否是安全的，例如教用户造炸弹，是不安全的

如上图所示，对 LLaMA-2 预训练模型经过对话用例微调之后，由人反馈参与的 RLHF 持续提升模型的能力，其打标流程主要有如下几个步骤：

标注者写一个 prompt
调用模型，随机参数，生成两个回答：response_1, response_2
标注者针对这两个回答标注如下信息：

a. 哪个回答更好：response_1 或者 response_2

b. 好的程度：非常显著的好，好一点，不确定（忽略不计的好）

c. response_1 是否安全

d. response_2 是否安全

e. 哪个回答更安全：response_1 或者 response_2

最终每一条标记数据组织成如下：

{
'prompt':P, # 如果是多轮对话数据，则除了最后一轮的回答，前文的问题和回答都在prompt里
'response_1':R1,
'response_2':R2,
'chosen_response':1 or 2,
'prefer_degree': 'significantly better', 'better', 'slightly better' or 'negligibly better/unsure',
'response_1_safe': 'yes' or 'no',
'response_2_safe': 'yes' or 'no',
'safer_response':1 or 2,

'打分人': xxx,
'response_1_model': model_name1,
'response_2_model': model_name2，
'标注时间': date,
'标注的超参': {
    '最长生成长度': max_len
    'temperature': temperature
    '多轮对话的超参': []
    ...
    }
}

这里，在调用模型生成两个回答时，LLaMA-2 会尽可能选择差异大的参数，使得生成的两个回答尽可能不同，方便用户进行准确的打标反馈。

写 Prompt 以及打分的思路

为了让 RLHF 训练能更有效的提升模型的能力，LLaMA-2 论文里提到了若干 Prompt 构造以及打分的 Tips。

对抗 Prompt 的编写

除了正常用户交互的 Prompt，为了产生更多具有代表性的数据，LLaMA-2 会让标注者写了一些有对抗性的 Prompt：

Risk category：可能会让 LLM 产生不安全内容的话题：
- 非法和刑事活动（例如恐怖主义、盗窃、人口贩运）
- 仇恨和有害活动（例如，诽谤、自残、饮食障碍、歧视）
- 不合格的建议（例如，医疗建议、财务建议、法律建议）
Attack vector：可能让 LLM 产生错误行为的提示：
- 心理操作（例如权威操作）
- 逻辑操作（例如错误前提）
- 句法操作（例如拼写错误）
- 语义操作（例如隐喻）
- 视角操作（例如角色扮演）

对抗 Prompt 的打分

对于上面写得对抗性 Prompt，模型需要按照下面的流程回答，才能算作好的回答：

首先解决安全问题
解释用户可能面临的潜在风险
在可能的情况下尽可能提供其他有用信息
避免负面的用户体验

负面的用户体验包括：

不要鼓励或者支持犯罪活动
不要鼓励或支持用户自己或者对他人进行危险行为
不要包含、鼓励或支持对用户或其他人的攻击性和辱骂行为
不要包含、鼓励或支持明确的性别相关内容

其他细节

1. 两个回答从什么模型来？

两个回答（Response）由不同的模型变体、不同的生成温度（Temperature）生成，目的是增加回答多样性。

但是需要注意的是，这里选择的模型，应该是正在训练的最新模型和上一版本的模型，以保证 Reward Model 是在最新的数据分布进行打分，给当前模型输出的打分足够准确。

论文里把这个过程称为 Iterative reward modeling data，是保证 Reward 模型在最新数据分布上使用当前版本模型和上一个迭代版本的模型进行采样的关键，同时还可以进行两个版本模型的评估比较。

2. 为什么只生成两个回答？

InstructGPT 是给定一个 Prompt 生成 K 个回答（K 通常为6左右），让标注者排序。与之对比，LLaMA-2 的标注方式有如下好处：

相同标注量下，Prompt 多样性更高；
同一个 Prompt，标注员的标注难度更低

3. LLaMA-2 数据量有多大？

LLaMA-2 每周更新一次标注数据。如下表所示，每一行是一周新标注的数据，LLaMA-2 总共约有百万条自己标注的数据：

表格含义如下：

batch：第几周
Num comparisons：有多少条数据，每一条数据如上文所述，包括了一个问题，两个回答，以及标注
Avg Turns：平均对话轮数
Avg Tokens：每条样本中的平均 Token 数量
Avg Tokens in prompt：Prompt 里的平均 Token 数量
Avg Tokens in response：Response 里的平均 Toekn 数量

总共 141w 条数据，但是每周都更换了生成回答的模型，这样来保证 Reward Model 学习的是微调后最新的生成结果。

4. 数据随时间的变化

由上表所示，每周标注的数据越来越多，这是因为标注者在不断的增加，他们也越来越熟练。同时对话轮数越来越多，增加了 Reward Model 学习的难度。

Prompt 越来越困难。LLaMA-2 论文里提到，刚开始要求标注者问一些简单的问题，后面要求标注者问更难的问题。下图可以看到，越往后的时间，Response 得到的分数越低，因为问题更难：

而 significantly better 的标签量越来越少，而 unsure 的标签量越来越多，这是因为模型的性能在不断的提升，所以采样到的两个回答都是高质量的。如下图所示，时间越往后，显著好的标签量越来越少，好的微乎其微的标签量越来越多：

5. 多轮数据怎么标注？

可以确定的是，多轮对话的数据体现在 Prompt 中，需要标注的是最后一轮的对话。但是构建多轮的 Prompt 论文没有详细提及，我们思考可以采用如下方法：

用户写一个 Prompt：P1

a. 单轮 Prompt：用户直接写

b. 多轮 Prompt：用户先和非标注模型对话多轮，最后一轮再进行标注
生成：内部调用两个模型，生成 response_1，response_2
标注：用户标注哪个回答是更优的，以及更优的程度，被选择的回答标记为 R1

这样标注数据的多样性得以保证，但是需要交互界面做一定的调整，让用户选择标注或者对话模式。

深入浅出 GPT-4 的体系结构

Vachel/小杞 — Thu, 27 Jul 2023 21:00:00 GMT

GPT-4（Generative Pre-trained Transformer 4）是 OpenAI 发布的最新 GPT 系列模型。它是一个大规模多模态模型，相比 GPT-3.5 / ChatGPT，GPT-4 可以接受图像和文本两种形式的输入，产生文本输出。输出依旧是一个自回归的单词预测任务。技术上，GPT-4 采用了专家混合（MoE）技术，进一步增强模型的能力。整体来说，GPT-4 在各种专业和学术基准上表现出了人类的水平，对于生成式的幻觉、安全问题均有较大的改善。

本文将和大家一起深入探索 GPT-4 的架构和能力。

论文地址：https://arxiv.org/abs/2303.08774

模型结构

GPT-4 的体系结构由16个不同的专家模型组成，每个模型都有 111B 个参数，总计约 1.76 万亿个参数。除了更大的参数规模外，另一个重要的细节是 GPT-4 使用了专家混合（Mixture of Experts，简写为 MoE）架构，这意味着模型中的不同组件或“专家”协同工作，每个组件都有助于最终输出。

上图可视了采用 MoE 的 Transformer 层。在 GPT-4 中 Attention 的参数量有 55B，MoE 的参数量是 111B * 16，一共 120 层 Transformer。每个 token 会通过一个路由算法选择两个 MLP 进行计算，参数 seq_len 为 8k，每个 MLP 分到 1k 个 token。

GPT-4 的模型宽度，深度基本和 GPT-3 (175B) 差不多，区别在于 MLP 的数量要多16倍。

并行策略

GPT-4 训练采用的并行策略是：

张量并行：8
流水并行：16
数据并行：196

总计使用约 3125 台机器（25000 张 A100）进行训练。其中 batch size 为 60M token，seq_len 为 8k。

张量并行和流水并行包含了 GPT-4 完整的模型参数，其结构如下：

其中，张量并行通讯耗时占比小于 15%，PipeDream 流水线气泡占比 28% 左右，Interleaved 1F1B 流水线气泡占比 16% 左右。

专家并行

专家混合（Mixture of Experts，简写为 MoE）是一种在模型增强的道路达到极限时采用的实用方法。运作原理很简单——由更简单的模型或“专家”组成一个小组，每个模型或专家都针对数据的特定方面，然后利用他们的集体智慧。

在实践中，这就像处理一项复杂的任务。比如建造一所房子会招聘建筑师、工程师和室内设计师等专业人士，每个人都有自己独特的专业知识。类似地，MoE模型由多个“专家”组成，每个专家都精通某个特定领域，协作以提供增强的结果。如下图所示：

对于 MoE 模型进行并行训练，这里可以做一个初步的估计：按照目前最大的 8 路张量并行，16 路流水线并行计算，每个 GPU 上的参数量是 14B，假设用 FP32 的梯度，则参数和梯度需要占用 84 GB 的显存，目前市面上还没有能完整放下显卡。因此，GPT-4 大概率会使用专家并行（Expert Parallelism）技术来节省显存占用。

模型能力

GPT-4 做的这些改进效果是显著的，相比前代 GPT-3.5/ChatGPT，他具备了新的能力：

具有图像理解能力，突破纯文字的模态，增加了图像模态的输入

如下图所示的例子，GPT-4 可以根据图标，计算格鲁吉亚和西亚的日均肉消耗量：
支持更长的上下文窗口

如之前所描述的，GPT-4 的 seq_len 为 8k，是 ChatGPT 上下文长度的2倍。
复杂任务处理能力大幅提升

在多类考试测验中以及与其他 LLM 的 benchmark 比较中发现，GPT-4 在更复杂、更细微的任务处理上，回答更可靠、更有创意。
改善幻觉、安全等局限性

LLM 的生成模式不可避免会有一些幻觉问题。GPT-4 在测试中发现其可以显著减轻在各类任务上幻觉问题，比最新的 GPT-3.5 模型提高 40% 左右。同样在安全能力的升级上，GPT-4 明显超出 ChatGPT 和 GPT3.5。

CC_Cleaner：一种丝滑高效且易扩展的数据清洗流程

WT/XT/MC — Thu, 06 Jul 2023 21:00:00 GMT

AGI 是数据 x 算法 x 算力的完美实践，是科研 + 工程 + 组织的优雅艺术。而在大模型训练的前置阶段，大数据的清洗是数据处理的基础。

以 Common Crawl 数据集为例，它可以轻松地在 Amazon 上获取到，是一个免费的、规模达到 PB 级别的网络爬虫数据集，其包括了超过 12 年间收集到的数据：原始网页数据（WARC）、元数据提取（WAT）和文本提取（WET）。如何对庞杂的原始数据进行清洗，得到高质量的数据集是一个值得研究的问题。本文将为大家介绍幻方基础研究实践优化后的数据清洗流程: cc_cleaner。

概述

FaceBook 在《CCNet: Extracting High Quality Monolingual Datasets from Web Crawl》工作中提出了从 WET 数据出发，清洗得到高质量数据集的流程。我们在此基础上做出了一些改进，提出了一种更为灵活、高效、强可拓展性的数据清洗流程，实现了更适合萤火集群的 cc_cleaner，其主要流程如下图所示：

相较于 cc_net，幻方 cc_cleaner 具有以下优点/特点：

更为完整的流程
- 相较于从 WET 出发的 cc_net，幻方 cc_cleaner 从原始的 WARC 文件出发，提供了更加完整的数据 pipeline
扩展性更高、流程配置化
- 数据清洗步骤中的算法、模型均可进行配置（如图所示），采用模块化方案，用户要自定义算法也只需继承原有函数，给定一个配置文件即可定义整个数据清洗流程
任务流程无感知
- 适配了萤火集群，任务被打断重启之后，不会重复处理数据，无需额外异常处理
- 自动进行任务分配，能方便地进行水平扩容
速度更快
- 基于高速文件系统 3fs（高达 7TB/s 读取，500GB/s 写入），大大提高了文件处理的效率

文本预处理

不同于 cc_net，我们的 cc_cleaner 从原始的 warc 数据开始做处理，通过 URL 过滤、文本抽取两步之后，生成对应的 WET 文件。其中 URL 过滤我们使用了一个包含 4.6M 个域名的黑名单，去除了包含欺诈/色情内容的网页。

清洗之后的 WET 文件在之后流程的使用中并不方便，为了更快、更方便地处理数据，我们为每一个 WET 文件生成了对应的 meta 文件，用于存储 WET 的信息摘要，其中包括：

原网页 URL
标题
对应 WET 文件路径
Data index（指明该 meta 对应的文件内容位置、长度等）
…

可以看到，meta 文件是 WET 的索引，通过更新 meta 中的内容，我们就可以实现对原始 WET 文件的修改、删除等操作。在生成 meta 文件时，我们会顺便生成文本的 minhash 用于去重，对于每一个文本，会生成一个 128 位的 int64 hash 值。

生成出来的 meta

去重

对于文本预处理中生成的每一个 meta 文件，有一个 128 位的 int64 hash 值，我们利用这个 hash 值进行去重。另外，cc_cleaner 的预处理内置了按行去重的功能，开启后将在处理 meta 时就把重复的行提前去掉。

对于 url page 的去重，首先按照去重率要求对 128 位 hash 进行分块，对于每一块 hash 的去重都是一个独立的任务，可以并行进行操作。cc_cleaner 会将每一块 hash 按照资源总量再次进行分配（例如 10 节点 * 256 核，就模 2560 进行二次分片），从前向后遍历数据，hash 相同的数据认为是重复的，并记录重复数据 index。之后将每块中记录的重复 index 从 meta 中删除，生成新的 meta，就实现了全局的去重。

语言分类

语言分类是对 WET 数据使用语言分类器进行分类，cc_cleaner 内置了 cc_net 使用的 fastText 分类器，用户也可以指定其他的分类器。通过遍历 meta 可以读取到原始的网页内容，对网页内容进行 lid 打分后，保留得分最高且大于 0.5 的 3 个分类，记录进 meta。如果没有语言得分超过 0.5，那这个数据将会被丢弃。

质量清洗

在抽取文本数据、去重并确定了网页数据的语言之后，仍然需要对文本进行一定的清洗来提高质量，称为质量清洗。幻方 cc_cleaner 内置了 cc_net 使用的 perplexity，这是一个在高质量数据上训练好的打分器，用来对 tokenize 之后的数据进行困惑度打分。

lid+困惑度打分示例

除此之外，cc_cleaner 还提供了以下清洗流程：

Document-level 过滤规则，例如：
- 丢弃单词数不在 50~100000 之间的文档
- 丢弃平均单词长度不在 3~10 之间的文档
- 丢弃“符号-单词”比例大于 0.1 的（’#’ 和 ’…’ 符号）文档
- 丢弃超过90%的行以 ‘•’ 开头或超过 30% 的行以 ’…’ 结尾的文档
- …
line-level 过滤规则，例如：
- 丢弃主要包含大写字母的行
- 丢弃只包含数字的行
- 如果行的内容是一个计数（例如：转发 12次评论 5条点赞 30个）就丢弃
- 丢弃只包含一个单词的行
- …

小结

通过上述流程，我们就生成了可以最终的 meta 数据，只需要通过 meta 数据读取原本的 wet 数据，进行 tokenize 之后就可以生成用于大语言模型训练的数据。可以发现，我们除了生成 wet 数据的文本提取步骤，其余的步骤都是对远小于原始数据的 meta 文件进行操作，用到的分类器、语言模型等均可以进行配置，这大大降低了研究人员修改数据清洗流程生成数据集的难度、时间。

HAI-LLM：高效且轻量的大模型训练工具

High-flyer — Tue, 27 Jun 2023 11:00:00 GMT

为了更好地发挥 GPU 集群的计算能力，训练具有惊人功能的、强大的万亿参数模型，一个高效简洁的大模型训练工具十分必要。幻方-基础研究最近研发了一款深度学习训练工具，名为 HAI-LLM，其实现了四种并行训练方式：ZeRO 支持的数据并行、流水线并行、张量切片模型并行和序列并行。这种并行能力适应了不同工作负载的需求，可以支持数万亿规模的超大模型并扩展到数千个 GPU。基于萤火集群的特性而自研的 haiscale 高性能算子，可以帮助 HAI-LLM 极大优化大模型训练的显存效率和计算效率。

本文将为大家介绍 HAI-LLM 工具的背后技术和使用方式。

显存与计算效率

首先，我们先了解训练巨大模型的显存和计算效率的挑战。

显存效率

训练万亿参数模型所需的显存远远超出了单张 GPU 的显存大小。在使用 Adam 优化器进行混合精度训练时，存储模型状态量（参数、梯度和优化器状态量）就需要约 12TB 的显存。作为比较，使用有 40 GB 的显存的英伟达 A100 GPU 训练大模型，仅仅为了存储模型状态，就需要 400 张这样的 GPU。

模型前向传播过程中会保存中间激活值，直到反向传播计算完 Loss 后才释放，这种额外消耗的显存会随 batch 大小而增加。假设 batch 设置为 1 的情况下，训练万亿参数模型就会产生超过 400 GiB 的激活显存。即使可以用 Checkpoint 方式（计算换显存）提前处理掉激活显存，但是对于训练而言，显存需求仍然非常之高。

必须在多个 GPU 设备之间有效地划分模型状态量和激活显存，才能让这种大模型在不耗尽显存的情况下开始训练。

计算效率

经估算端到端训练一个万亿参数的模型大约需要 5000 Zflops（即 5 后面带有 24 个零；这个估算结果基于 OpenAI 的研究《law of scaling》。这意味着训练这样一个模型需要 4000 张 A100 以 50％的计算效率运行大约 100 天。

尽管大型超级计算 GPU 集群可以拥有超过 4000 个 GPU，但是由于 batch 大小的限制，要在这种规模上实现高计算效率仍然是一项挑战。计算效率随着计算时间对通信时间的比例的增加而增加。该比例与 batch 大小成正比。但是，训练模型的 batch 大小有一个上限：超过这个上限收敛情况会明显变差。

实际上最大的模型之一，GPT-3 的训练 batch 大小约 1500。如果使用大约 4000 张 GPU，即使我们可以自由设置 batch 大小为 4000，每张卡上的 batch 大小也只有 1，这将影响扩展性。

3D并行

设计高效的并行方式对训练大模型至关重要。HAI-LLM 结合萤火集群的特性，可以实现三种并行训练方式的灵活组合：数据并行、流水并行、张量并行，并使用序列并行方式对 Transformer 进一步并行优化，极大地提升显存利用和计算效果。

彩色块显示有 32 张 GPU （八个计算节点，每个节点有四个 GPU）的系统上进行 3D 并行训练的 GPU 映射。同一颜色的 GPU 在同一节点上

数据并行

为了减少显存的占用，HAI-LLM 采用 ZeRO 数据并行的方式，把训练优化器的状态均分到 1/N 个 GPU 上（N 是数据并行的数量），如下图所示：

在做前向和反向传播时，ZeRO 数据并行会先做 allgather 获得完整的参数，然后在前向和反向传播结束后释放掉，只保留 1/N 的参数和梯度。

使用比较简单，HAI-LLM 提供了hai_llm.optim.ZeroOptimizer 工具，该工具可以自动切分优化器状态并执行梯度同步。代码示例如下：

from hai_llm.optim import AdamW, ZeroOptimizer

opt_cls = partial(AdamW, lr=1e-3, weight_decay=0.1)
optimizer = ZeroOptimizer(
    model.parameters(),
    opt_cls,
    group=dp_group,
    max_grad_norm=cfg.optim.max_grad_norm,
    grad_scaler=scaler
)

训练的时候只需要执行 optimizer.step() 就可以完成数据并行的梯度同步和参数更新。

流水并行

流水并行把模型切分成 M 个阶段，分别放在 M 个 GPU 上，能够大幅节省显存占用和提升扩展性能。训练的时候每个阶段接收上一阶段的输出作为输入，然后传给下一个阶段。如下图所示：

可以看到，这种并行方式会造成大量的 GPU 空闲（Idle）。为了减少浪费，Gpipe 和 PipeDream 的概念被提出：如果同时进行多个迭代，每个节点在同一时刻负责不同的迭代的计算，就可以避免数据依赖，不用在原地干等了：

在 HAI-LLM 中，模型的切分是在 Builder 里由用户切分的（比如 GPTBuilder），用户可以继承 hai_llm.builders.base_builder.BaseBuilder 并实现自己的 build_model 方法，在该方法中创建每个 GPU 对应的模型实例并传入 _blocks2model 中。Builder 创建出来的模型包含一个 forward_backward 方法，用户可以传进输入、损失函数、标签，该方法会执行前向传播和反向传播并返回 Loss：

loss, _ = gpt.forward_backward(x, criterion=criterion, labels=(y,))

上面等同于以下过程：

output = gpt(x)
loss = criterion(output, y)
loss.backward()

张量并行

张量并行把全连接层的参数和计算切分到 K 个 GPU 上，能够节省显存占用和提升扩展性能。假设 K=2，张量并行把矩阵乘法 X @ A @ B 变成 X @ A1 @ B1 + X @ A2 @ B2，其中 A 沿着纵轴切成 A1 和 A2，B 沿着横轴切成 B1 和 B2；GPU 0 会存放参数 A1 和 B1 并负责 X @ A1 @ B1 的计算，GPU 1 则存放参数 A2 和 B2 并负责 X @ A2 @ B2 的计算。

HAI-LLM 中的模型均支持张量并行的方式，如果用户需要自定义模型结构，我们在 haiscale.tensor 中提供了以下四个用来构建张量并行模型的工具：

ColumnParallelLinear：沿着纵轴切分的全连接层算子；
RowParallelLinear：沿着横轴切分的全连接层算子；
VocabParallelEmbedding：沿着字典维度切分的 Embedding；
vocab_parallel_cross_entropy：接收沿着字典维度切分的输入，进行交叉熵损失函数计算。

用户使用 haiscale.tensor 构建模型可以不需要关心具体的张量并行通讯的实现。具体的使用方法可以参考 GPT 和 LLaMA 的模型代码。

序列并行

Megatron-LM 提出的序列并行方案旨在通过一些方式分摊张量并行中无法分摊的显存。

如上图所示，在张量并行的基础上，将 Transformer 核的 LayerNorm 以及 Dropout 层的输入按序列长度维度进行了切分，使得每个 GPU 上面只需要做一部分的 Dropout 和 LayerNorm 。这样做进一步降低了计算资源和显存的开销。

HAI-LLM 实现了 allreduce_sequence_parallel 方法，用户可以通过一行代码实现模型中 Transformer 算子的序列并行。如下代码所示：

cfg.gpt.sequence_parallel = True

for step in range(resume_step, cfg.train.steps):
    ...

    allreduce_sequence_parallel(gpt.module)

    ...

训练配置

HAI-LLM 通过 Python 文件进行模型图纸、超参数配置，通过 load_config 函数加载后会把字典转换成 AttrDict 对象，能够直接通过属性访问。使用方式如下例所示：

from hai_llm.utils import load_config

cfg = load_config(os.environ.get('CONFIG', 'toy/toy_config.py'))
model = cfg.model_builder.cls(cfg).build_model()

配置文件中至少需要描述清楚如下几个部分：3D 并行设置 parallel，模型图纸 model_builder，超参数 model_param 和中间结果保存 checkpoint。HAI-LLM 目前提供了 GPT-3，LLaMA，PaLM 三种模型的复现和配置案例。

效果评估

HAI-LLM 提供了 /eval_benchmark/eval.py 脚本在验证集上测试模型的 Loss，其需要通过 hfai 提交到集群中运行。输入模型的配置文件、Checkpoint 位置、验证集数据等信息，指定对应的评估方法即可。如下所示：

hfai python eval_benchmark/eval.py \
    --cfg <config-file> \
    --ckpt_path <checkpoint-dir> \
    --load_step <load_step> \
    --val_ckpt_path <val-checkpoint-dir> \
    --data_path <val-data-path> \
    -- -n xx -p 40 -f -i ubuntu2004-cu113

同时，/eval_benchmark/ 下也提供了几种常用的模型效果评估方法，方便用户对训练好的模型进行评测。评估方法包括：

hellaswag
winogrande
piqa
triviaqa
mmlu
race

HAI Platform：开源的 AI 训练平台

High-Flyer — Thu, 11 May 2023 21:00:00 GMT

当下 AI 大模型如火如荼的发展，许多厂商开始构建起基于 GPU 的大规模深度学习训练集群。随着算力需求的扩大，一个能对 GPU 资源进行统一高效利用的 AI 平台，越来越成为各 AI 团队降本提效的重要武器。

构建一个能管理大规模 GPU 集群资源的 AI 平台，主要会遇到如下几个痛点问题：

资源调度：算力规模不断扩大，而训练任务的计算需求又多种多样，如何处理任务和算力的关系以最大化集群资源使用？
使用效率：集群使用有峰谷差异，同时又要应对突发任务的需求，如何兼顾时效性和集群整体效率？
迭代适配：集群会有升级迭代，在调整算力规模、类型、使用规则等场景下，如何让平台快速适配，尽可能降低切换成本？

针对上述问题，幻方 AI 在自建的萤火集群上进行多年的研发与测试，积累了一套高效管理 GPU 集群资源的 AI 训练平台方案，名叫 HAI Platform。其以任务级分时调度共享 AI 算力的理念将集群零散资源进行整合再分配，成功支持在 1500+ 计算节点上稳定运行深度学习训练和其他多类型任务，日常算力占用率 95% 以上，日常 GPU 使用率 75% 以上，计算和存储节点间的数据吞吐 7TB/s 以上，服务了上百个 AI 团队。

现在，幻方 AI 将正式开源 HAI Platform，将这套久经考验的 AI 训练平台服务于更广阔的 AI 社区。HAI Platform 可以被部署在私有集群或公有云的 GPU 资源中，帮助您高效利用 GPU 集群资源，提升团队整体研发效率。我们希望让更多“想象力”和“创造力”生长，期待与各方科学家及开发者们一同共建 AI 时代。

产品地址：https://github.com/HFAiLab/hai-platform

产品文档：https://hfailab.github.io/hai-platform/

安装与设置

使用 HAI Platform，首先需要如下基础设施：

一个部署 K8S 的算力集群，包括管理节点、计算节点、服务节点（用于服务部署、数据库、监控等）；
一个或多个存储集群，让所有计算节点可访问的文件系统，存储信息包括代码，数据，日志，配置文件等；
算力集群与存储集群之间所有节点通过高速网络互联，建议使用 RDMA 。

接下来，您可以按照开源仓库中 HAI Platform 搭建指南进行安装和部署。

任务调度

任务是 HAI Platform 使用的基本单位，而非用户。所有用户提交的任务都由 HAI Platform 统一管理和调度。

用户独占 vs 分时调度

幻方 AI 提出分时调度理念对集群资源进行管理。用户提交任务，如运行 python / bash 代码，启动开发容器等，由平台根据当前资源需求、集群忙闲程度等进行任务的中断和加载。任务代码需要遵循平台编码规则以确保可以断点续跑，具体流程如下：

接受集群的打断信号；
保存 checkpoint（模型参数，优化器参数等）；
通知集群打断；
从 checkpoint 恢复，继续运行。

注意！ 部署 HAI Platform 的集群不会将 GPU 资源池化，而是以计算节点为基本单位，根据资源类型、网络区域等条件进行分类标记。HAI Platform 鼓励用户一次性用满多张 GPU，进行并行训练。用户提交任务时需选定节点数量 n，则该任务可获得 n 的整数倍个 GPU，比如在8卡计算节点上提交使用4个节点的任务，则该任务会获得32张 GPU 进行并行训练。

更详细的内容可以阅读《分时调度》。

用户管理

HAI Platform 提供以配额的方式记录集群的各类资源，通过优先级管理用户的使用权限。例如：某用户在 NORMAL 优先级上拥有10个计算节点的配额，其在 HAI Platform 上提交若干个任务，则最多同时以 NORMAL 优先级调度10个计算节点运行。算力资源优先满足高优先级用户的训练需求，同优先级内则交替使用集群算力资源。

任务优先级列表

同时，每个用户都可以归属于某一个用户组。计算资源和优先级配额是按用户为单位进行分配的，而同组的用户可以共享自定义环境、私有数据存储。

此外，HAI Platform 还提供了管理员用户身份，其有权限执行用户创建、停用、归档等操作，调整各类资源和优先级配额，管理后台开发容器等。

环境管理

部署 HAI Platform 的集群会包含基础环境和用户自定义环境两种环境管理模式：

基础环境由平台统一管理，包含大部分主流依赖包，任务运行时默认加载。通常根据系统、GPU 驱动及 CUDA 版本等要求统一管理若干个基础环境。基础环境供平台所有用户共享；
自定义环境由用户自行安装依赖，可以基于基础环境进行扩展。自定义环境可以在同组内共享。

import haienv
haienv.set_env('env_name')

两种环境都可以通过 haienv 工具进行管理，适用于各种任务场景。通过基础环境的构建，我们希望尽可能降低 AI 研发前期的环境构建成本。

更详细的内容可以阅读《环境配置》。

可视化交互

为了方便用户与 GPU 集群的交互，HAI Platform 提供了用户界面，HAI Platform Studio ，其主要功能包括但不限于：

管理和启动基于 JupyterLab 的开发容器，配合 Jupyter 插件，可以进行代码调试或任务提交；
管理提交的任务，支持任务的实时监控、日志分析和任务状态管理；
集成了讨论区、任务性能监控、任务调度管理等功能。

另外，HAI Platform Studio 子应用 /monitor 集成了一定的监控和管理能力，包括：

用户及任务管理：创建用户；修改用户分组、Quota；管理用户开发容器和任务；
节点管理：查看分组节点以及各个节点的详细信息。

HAI Platform Studio 概览界面

更多内容可以阅读开源仓库：https://github.com/HFAiLab/hai-platform-studio。

数据管理

部署 HAI Platform 的集群会将训练数据存储在文件系统中。训练中的任务会实时从文件系统中读取数据，流转到计算节点。这里，优良的文件系统将是提高 GPU 利用率，进而提升集群整体效率的决定性因素之一。

为了最大化 GPU 利用率，降低数据读取的成本，幻方 AI 自研了高速文件系统 3FS，其专门针对深度学习训练中样本读取的场景进行优化，能支持计算和存储节点间 7TB/s 以上的数据吞吐。3FS 需要用户将原始数据中较多的小文件聚合成大文件进行存储，推荐使用 FFRecord 格式规范来聚合文件。

GTC 2023 | 按需分配的AI算力 Ⅱ

High-Flyer — Mon, 03 Apr 2023 21:00:00 GMT

“切勿错过AI的决定性时刻”这句 NVIDIA GTC 的广告语自周三起，一时间传遍全网。依照惯例，大会由熟悉的面孔黄仁勋围绕AI、芯片等科技，发布了一系列前沿技术产品。

今年幻方 AI 再度受邀，在 NVIDIA GTC 2023 大会中进行了一次技术主题分享。

自 2019 年起，我们为了满足自己科研作业的大规模算力需求，逐步构建了幻方萤火深度学习智算平台，其包括存储网络相关基础设施、分时调度算法、以及深度优化的 hfai 框架。在实际的工作中，我们常常遇到很多模型训练与落地的挑战，本次讲座分享的内容就是过去一段时间，我们在软件和硬件上，取得的一些突破以及积累的相关解决方案。

以下是本次 GTC 主题分享的内容概要，也是我们阶段性的最佳实践。

FFRecord

目前网上的开源文件格式，比如 TFRecord 等在不需要随机访问样本的情况下能够获得优秀的性能；但在很多机器学习的场景下，用户训练模型时每一步需要随机采样多个样本。为此，我们专门设计了 FFRecord (Fire Flyer Record) 文件格式，帮助用户更加方便地享受 3FS 高速存储批量读取带来的效率提升。同时我们实现了和 PyTorch 完全兼容的接口，方便用户的迁移和加速。

用 ImageNet 性能测试显示，用 FFRecord 能够比大量小文件读取的 Pytorch dataloader 性能提升近 40%。不仅如此，为了方便用户在萤火集群上的训练，我们构建以 FFRecord 为基础的幻方 AI 数据集仓库。目前数据集仓库里已收纳了 20 多种大规模公开数据集，近 120 多个用户私有数据集。

hfai.nn

幻方 AI 对 Pytorch 框架进行了深度优化，结合幻方萤火集群的特点，对一些常用的 AI 算子重新研发，提升效率，进一步提升了模型整体的训练效率。

更快。为了让模型计算得更快，我们通过优化 CUDA kernel，融合多运算，优化模型结构等，使得深度学习算子计算更高效。

更省。现代模型越来越大，在不改变精度的情况下，我们通过压位等技术，在训练时减少了部分常用函数中间变量的内存，大幅节省训练的内存占用。

haiscale

随着模型规模和数据规模的扩大，越来越多的模型训练都要在多机多卡上进行。因此并行训练中的优化至关重要。Haiscale 是一个轻量级的高性能并行训练工具，集成了幻方 AI 多年的并行训练研发经验，能很轻松地兼容 PyTorch 进行并行训练提速，其包含：haiscale.ddp, haiscale.fsdp, haiscale.pipline 三大套件，它们都以幻方自研的 hfreduce 多机多卡通信方案为底座，进行并行训练的加速。

haiprof

haiprof (High-flyer’s AI Profile) 是一个针对 PyTorch 模型的性能分析工具，用户能够在无需修改代码的情况下分析模型性能瓶颈，能随时调用，为整个训练过程做全方位的“CT”。

在任务运行的过程中，haiprof 会收集以下信息：

CPU/GPU 利用率、功率、内存占用
用户读写流量、IB 流量
内存带宽、跨 NUMA 带宽
forward、backward、optimizer、dataloader 的耗时

haiprof 还支持对 PyTorch 代码进行更深入的 tracing，分析不同函数、算子的时间开销。

自研路由算法

对于影响深度学习训练快慢的因素，人们常常容易忽略网络传输在训练提速中的重要作用。特别是在大规模集群，分布式训练的场景中，网络的拥塞可能直接导致 GPU 算力的失效，就像空有一段段双向 8 车道的快速路，但如果道路规划凌乱，高速路也只能沦为大型停车场。

在实际的应用中，我们便发现网络拥塞无处不在。而自研路由算法，帮助我们突破了传统的 fattree 架构，带我们来到了更加广阔的世界。

这一算法，能够支持我们多台大型 800 口互联，集群算力的提升突破了物理上传统 fattree 的上限；
支持我们使用同一个网络来支持存储和互联，大大降低了集群搭建的成本，提升了可维护性，增加了稳定性；
我们能够在原有典型 fattree上，构建其他低速子树，支持其他的异构运算节点，让所有节点都有各自需要的互联能力，提高了平台算力的多样性。

haiprof | 模型训练性能分析工具

小杞 — Mon, 27 Mar 2023 21:00:00 GMT

全世界都在给 GPT 的训练和推理算账，以至于微软会迫不及待地分享“每美元的推理能获得 2 倍的性能”的喜讯，其传递的信号或许是：一丁点的算力浪费，都在这场 AI 竞争里是致命的。

而越往后发展，模型训练效率的高低，越发成为科研人员不容忽视的难点。无论从加快训练测试更多参数，还是从节省研发成本的角度来说，低效的任务训练已经变得不可容忍。摆在眼前的，对于科研用户来说，一套趁手易用的模型性能分析工具，是工作中的必备利器。那么如何快速得给模型做一个全方位的“体检”，如何更简单高效的地定位模型训练中的“病灶”？今天就为大家介绍幻方 AI 研发的模型性能分析工具 haiprof。

haiprof

haiprof (High-flyer’s AI Profile) 是一个针对 PyTorch 模型的性能分析工具，用户能够在无需修改代码的情况下分析模型性能瓶颈。在任务运行的过程中，haiprof 会收集以下信息：

CPU/GPU 利用率、功率、内存占用
用户读写流量、IB 流量
内存带宽、跨 NUMA 带宽
forward、backward、optimizer、dataloader 的耗时

haiprof 还支持对 PyTorch 代码进行更深入的 tracing，分析不同函数、算子的时间开销。

使用方式如下：

hai python train.py -- --nodes 8 --options profile.time=120 --options profile.trace=1
haiprof train.py --time 120

先提交任务，后开启分析工具。接下来我们将会以一个真实场景下，用户模型训练的优化过程作为示例，来演示 haiprof 的使用方法。

模型运行全景 profile 请求

1. 查看任务Log

这是一个正在运行的训练任务的日志，在日志中我们打印出了每个 step 的 forward、backward 耗时。我们可以看到，这个任务的 forward 耗时要远远大于 backward 。通常来说，大部分模型的 forward 时间和 backward 是差不多或者更小的。这个训练任务显然存在着一些问题。接下来我们用 haiprof 分析一下该模型。

2. 使用 haiprof 命令行工具

这里我们先使用 hfai prof 命令工具对正在运行中的任务做一个简单的 profile，使用上只需要指定任务的 ID 或者名字，然后设置 profile 的时长。

haiprof train.py –time 60

现在我们成功提交了对 train.py 这个任务的 profile 请求，然后我们需要等待 60 秒，profile 完成后会把结果写入 log 文件夹。接下来我们可以使用 haiprof 提供的可视化工具进行分析。

3. 运行数据可视化分析

我们在 haiprof 库提供了一个 Analyzer 工具，用户只需要输入 profile 的 log 文件夹路径即可。

我们先调用 show_overview 函数来看看任务的整体运行情况

从上图中我们可以看到整个任务的 IB 平均流量、 CPU、GPU 平均利用率，以及内存、显存的占用情况等等。我们观察到这个训练任务的 GPU 利用率非常低，只有 30%，而 CPU 利用率却有 60%~70%。

接下来我们通过 show_one_rank 函数来看看详细的情况

这里我们看到 GPU 的利用率时高时低，有一半的时间的利用率都接近 0。

由此我们可以推测该训练任务阻塞在 CPU 运算上。

任务毫秒级 tracing

1. 使用 haiprof 命令行工具

为了更详细地分析模型的运行时间，我们可以使用 haiprof 的 tracing 功能，用户只需要在使用 hai python 提交任务的时候指定选项 profile.trace=1，就可以自动地对任务进行毫秒级别的 tracing。

hai python train.py -n 1 -p 50-f –opons profile.trace=1 –options profile.warmup=120 –options profile.time=15

该命令代表先预热 120 秒，然后 tracing 15 秒

2. 运行数据可视化分析

Tracing 结束后结果以 json 格式保存在 log 目录里。我们可以把结果拷贝到当前目录然后使用 tensorboard 查看结果：

打开 tensorboard，在 overview 页面我们可以看到模型运行过程中的 GPU SM 利用率、不同部分的时间占比等。

在 GPU kernel 页面我们可以看到不同 GPU kernel 算子的时间占比。

在 trace 页面，我们可以看到整个模型运行的 timeline，里面包括的 forward、backward、optimizer 的时间，以及不同算子、函数的时间占比。

这里我们看到，这个模型中有相当一部分的算子运算是在 CPU 上做的，这导致了在大部分的时间 GPU 都处于空闲状态。

诊断后优化

接下来我们根据这个发现优化一下模型，把这部分算子放到 GPU 上计算。改完后重新提交任务：

hai python train.py -n 1 -p 50 -f –name v2

我们可以看到，优化过后的模型 forward 的时间只有原来的 1/3，整体训练速度提升了一倍。

全球气象预测开源大模型 OpenCastKit

Vachel — Fri, 17 Mar 2023 21:00:00 GMT

用 AI 方法对现代数值天气预报（numerical weather prediction, NWP）进行改进提升近两年受到了广泛关注，如 Nvidia 发布的 FourCastNet、 DeepMind 发布的 GraphCast 和华为发布的盘古气象大模型，在与欧洲中期天气预报中心（ECMWF）的高分辨率综合预测系统（IFS）对比中都获得了不错的效果。

基于此，我们最近复现整合了这些工作，并将这些成果贡献给开源社区。我们基于 FourCastNet 和 GraphCast 论文构建了一个新的全球AI气象预测项目 —— OpenCastKit。

项目地址：https://github.com/HFAiLab/OpenCastKit

这个项目提供了一个强大的、基于ERA5数据训练的开源气象模型和参数，可以生成全球高分辨率的气象预测。具体来说，它包含：

一个统一的数据处理工具，抽取ERA5数据和特征并整理成高性能训练数据格式 FFRecord；
基于 hfai 算子和 hfreduce 并行通信优化的 FourCastNet 模型源码和 GraphCast 模型源码，供社区研究优化；
基于1979年到2022年15TB的ERA5数据，在萤火高性能集群上训练的模型参数，可以进行微调，获得高精度预测结果。

同时，我们上线了一个每日更新的 HF-Earth，实时展示气象大模型输出的全球预报效果：

Demo地址：https://www.high-flyer.cn/hf-earth/

经过一段时间的测试来看，AI气象大模型对台风、极端降水等事件的预测上效果明显，在长期气候变化的分析中可以起到一定作用。希望在此开源项目的基础上，构建出更加强大的 AI 气象应用。

数据集

欧洲中期天气预报中心（ECMWF）提供了一个公开可用的综合数据集 ERA5，其将物理模型数据与来自世界各地的观测数据结合起来，形成一个全球完整的、一致的数据集，以小时级到天级不等，提供包括温度、风量、降水、水文、气压等多项全球气象指标数据，供各种气象预报模型学习。

官方地址：https://www.ecmwf.int/en/forecasts/datasets/reanalysis-datasets/era5

FourCastNet 与 GraphCast 使用了不同规模的 ERA5 数据来训练，产生的预测效果各自不同。前者仅使用了 20 个相关气象指标，包括 4 个不同位势高度下的温度、风速、相对湿度和一些近地表变量，其旨对极端天气、自然灾害进行预警；而后者使用了更加全面的数据，其包含 37 个不同位势高度下的气象指标和 5 个地表气象指标，总计 227 个指标，其旨在对气象变化进行更加全面的评估和预报。

对此，我们将这些数据进行了归纳整理，通过 hfai.datasets 工具进行管理优化。原始数据通过特征处理，转化成 “ $X_{t-1},X_t \rightarrow X_{t+1}$ ” 的模式，通过高性能训练样本格式 ffrecord 进行保存，从而可以在萤火集群中进行高效的并行训练。更多信息可以浏览 hfai 数据集仓库。

模型构建和优化

为了进行 0.25° 分辨率下的全球气象预测，FourCastNet 采用自适应傅里叶神经算子 AFNO，而 GraphCast 采用了图神经网络。前者计算效率高效，可以灵活且可扩展地建模跨空间和不同指标之间的依赖关系；后者通过构建节点之间的联系，更加详细捕捉如“蝴蝶效应”般的气候因子影响。前者在小 batchsize 上可以进行数据并行以加速训练，而后者球体节点之间的 message passing 参数规模更大，需要进行流水线并行（或称模型并行）的改造，以实现模型的完整训练。

这里我们采用自研的 haiscale 高性能并行训练工具库对两种模型进行复现优化。对于 FourCastNet，我们使用 haiscale.ddp 或者 haiscale.fsdp 进行数据并行优化，实验中我们采用小 batchsize 即实现了论文效果的复现；对于 GraphCast，完整参数基本无法塞入单张显卡，因此对于不同的环节，如球体中 grid 节点与 mesh 节点进行message passing，我们需要对其拆分，让其分布在不同的显卡上，通过 haiscale.pipeline 对不同环节进行串联，实现模型并行训练。具体如下：

FourCastNet 数据并行

FourCastNet 模型的训练包括 pretrian、finetune 和 precipitation 三个部分。模型采用递进式，即以 $X_t$ 作为输入，预测下一步 $X_{t+1}$ 。一次训练输出多步，与真值对比计算 loss。如下伪代码所示：

from hfai.datasets import ERA5
from haiscale.ddp import DistributedDataParallel
from torch.utils.data.distributed import DistributedSampler

model = FourCastNet(args).cuda()
model = DistributedDataParallel(model)
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)

data = ERA5(split='train')
sampler = DistributedSamper(data, shuffle=True)
dataloader = data.loader(args.batch_szie, sampler=sampler, num_workers=8, pin_memory=True, drop_last=True)

# training ...
for step, (xt0, xt1, xt2, pt2) in enumerate(dataloader):
    xt1_pred = model(xt0)            # pretrain
    xt2_pred = model(xt1_pred)          # finetune
    pt2_pred = model(xt2_pred, precip=True)    # preciptation

    pretrain_loss = criterion(xt1_pred, xt1)
    finttune_loss = criterion(xt2_pred, xt2)
    precip_loss = criterion(pt2_pred, pt2)
    
    # optim ...

# stop hfreduce
model.reducer.stop()

haiscale.ddp 默认采用 hfreduce 进行通信优化，我们还可以使用优化算子，加入一行 model = hfai.nn.to_hfai(model) 代码进行进一步加速。在萤火集群上我们使用 96 张 A100 进行数据并行加速，耗时 16~17 个小时左右基本可以完成 FourCastNet 的整体训练。

GraphCast 数据并行

不同于 FourCastNet，GraphCast 只有主干模型一个，其也是采用递进式，不过以 $X_{t-1}, X_t, T, C, G$ 作为输入，预测下一步 $X_{t+1}$ 。这里 $T, C$ 代表了时间戳信息和地理位置信息， $G$ 代表所构建的球体 Graph 信息。如下伪代码所示：

from hfai.datasets import ERA5
from haiscale.ddp import DistributedDataParallel
from haiscale.pipeline import PipeDream, make_subgroups, partition
from torch.utils.data.distributed import DistributedSampler

dist.init_process_group(...)
torch.cuda.set_device(local_rank)
rank, world_size = dist.get_rank(), dist.get_world_size()

dp_group, pp_group = make_subgroups(pp_size=pp_size)
dp_rank, dp_size = dp_group.rank(), dp_group.size()
pp_rank, pp_size = pp_group.rank(), pp_group.size()

model = GraphCast_sequentail(args)
model = partition(model, pp_group.rank(), pp_group.size(), balance=[1,1,1,1,1,1,1,1])
model = DistributedDataParallel(model.cuda(), process_group=dp_group)
model = PipeDream(model, args.chunks, process_group=pp_group)
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)

data = ERA5(split='train')
sampler = DistributedSamper(data, num_replicas=dp_size, rank=dp_rank, shuffle=True)
dataloader = data.loader(args.batch_szie, sampler=sampler, num_workers=8, pin_memory=True, drop_last=True)
earth_graph = generate_graph(args)

# training ...
for step, (xt0, xt1, xt2) in enumerate(dataloader):
    loss = model.forward_backward(xt0, xt1, earth_graph, criterion=criterion, labels=(xt2,))
    
    # optim ...

# synchronize all processes
model.module.reducer.stop()
dist.barrier()

在使用 haiscale.pipeline 进行流水线并行训练时，需要我们提前将模型进行拆分，通过 haiscale.SequentialModel 进行模型的串联。同时 haiscale 提供了一个统一的 forward_backward 接口，进行样本和标签的统一输入和结果输出。在萤火集群上我们使用 256 张 A100 进行模型并行加速（单节点 8 卡做流水并行，32 节点做数据并行），耗时 3 天左右基本可以完成 GraphCast 的整体训练。

关于代码的更多细节可以访问项目地址阅读源码。

预测结果

参照论文中的评估方式，我们采用递归输出未来多天的预测结果，与真实值对比，通过误差增长曲线来比较不同AI气象大模型的预测效果。如下图所示：

可以看到，在进行 14 天的中期天气预报测试中，无论是GraphCast 还是 FourCastNet，递归预测导致误差随时间逐步增长。整体误差上看 GraphCast 考虑了地理时间和地理位置的因素，预测误差比 FourCastNet 要小。受此启发，我们将时间和地理信息加入 FourCastNet 进行模型训练（FourCastNet+），发现最终模型输出的预测误差几乎与 GraphCast 一致。

下面我们以 2022 年 6 月 22 日开始连续输出 14 天的预测，展示 OpenCastKit 的预测效果：

FourCastNet 温度预测与风力预测

GraphCast 温度预测与风力预测

真实温度与风力

可以看到 FourCastNet 和 GraphCast 都可以对风力和温度的衍变进行比较准确的预测。其中 GraphCast 相对来说更加接近真实情况，包括气象指标的细节纹理更丰富和一致，还有在6月30号开始在我国东南沿海的两次台风路径的预测。

萤火跑模型 | 高性能 Stable Diffusion 助力高质量 AI 绘图

Daylight & Vachel — Thu, 22 Dec 2022 14:00:00 GMT

AI 绘画最近成功破圈，成了炙手可热的热门话题。DALLE，GLIDE，Stable Diffusion 等基于扩散机制的生成模型让 AI 作图发生质变，让人们看到了“AI 转成生产力”的曙光。

在这些扩散模型中，Stable Diffusion 以其优秀的效果和开源的权重成为了其中的代表，受到广泛的关注和体验。其基于 Laion5B 超大规模“文本 - 图像”对数据集，Stable AI 宣称用了 5000 张 A100 耗时几个月训练而成。幻方 AI 近期在萤火二号上使用 Google Caption 数据集复现了 Stable Diffusion 的训练，并进行了优化。通过幻方自研的 hfai.pl 插件将源代码 Pytorch Lightning 框架与萤火集群的特性轻松整合，并通过 3FS、hfreduce、算子等优化工具对模型训练提速。

本文将分享我们对 Stable Diffusion 训练优化的心得体验，帮助研究者和开发者们降低研究门槛。

论文标题：High-Resolution Image Synthesis with Latent Diffusion Models

原文地址：https://arxiv.org/abs/2112.10752

源码地址：https://github.com/CompVis/stable-diffusion

模型仓库：https://github.com/HFAiLab/stable-diffusion

模型介绍

Stable Diffusion 基于 Latent Diffusion 进行了扩大训练，其将 Text Encoder 从 BERT 更换为 CLIP Text Encoder。我们首先了解下 Latent Diffusion 的模型设计。

过往虽然扩散模型的生成能力非常强大，在许多不同类型的生成任务上都达到了 SOTA 的水平，但由于其迭代生成的特点在训练和推理时往往都需要耗费非常多的 GPU 资源。Latent Diffusion 针对这一点做出了改进，通过将扩散过程从图片的像素域转变为在编码后的潜空间上，这种方式大幅降低了扩散模型运行时的复杂度，同时也能保留较好的细节和图像生成效果。整体结构如下图所示：

Latent Encoding

Latent Diffusion 通过在一般的扩散模型中增加一个变分自编码器来实现像素空间和隐空间的转换。在训练 DDPM 模型前首先会在 ImageNet 上训练一个 VAE 模型，其旨在学习一个编码器和解码器，用于将图像压缩成潜在编码。VAE 会将图片压缩到原来的 1/4 ~ 1/8 的大小，因此可以极大的降低在隐空间进行操作时的计算复杂度。而在这种情况下压缩再还原后的图片也能较好的保留原图片的信息，基本不会因为压缩产生过大的损失。

Cross-Attention for Multimodal Generation

Latent Diffusion 创新的使用了注意力机制来进行条件控制信息和图像生成过程的融合。在条件控制生成上，作者在 U-Net 模型的每一层都引入了条件控制信息来控制图像生成的方向。而在控制信息和图像生成过程的融合上，作者引入了 Cross Attention。基于注意力的融合机制使得模型可以更容易的使用不同类型的条件控制信息，例如文本生成图片，图片生成图片，或是语义图生成图片等等。

Text-based Generation

和 Latent Diffusion 不同，Stable Diffusion 的重点在于文本生成图片。Stable Diffusion 使用了远大于 Latent Diffusion 的 LAION-5B 数据集中的 25 亿“图片 - 文本”对进行训练。此外受到 Imagen 等其他生成工作的启发，还将 Latent Diffusion 中使用到的 BERT Text Encoder 替换成了效果更好的由 CLIP/ViT-L-14 预训练的 Text Encoder。

模型实践

训练数据集

为了验证 Stable Diffusion 模型的训练性能，我们使用 Google Conceptual Caption 数据集复现了 Stable Diffusion 的训练。Google Conceptual Caption 是一个相对小范围的多模态数据集，其中有 285 万“图像 - 文本”对。该数据集已集成在幻方 AI 的数据集仓库中，转化为 ffrecord 训练数据格式存储在 3FS 高速存储里。用户可以通过如下方式调取获得高速的训练数据读取：

from hfai.datasets import GoogleConceptualCaption
dataset = GoogleConceptualCaption(split="train", transform=transform)
dataloader = dataset.loader(**args)

hfai.pl

Pytorch Lightning (PL) 在 PyTorch 基础上进行了封装，拥有其独特的并行训练接口。Stable Diffusion 源码基于 PL 所构建，为了使其利用起萤火集群的各种优化特性，我们采用幻方开发的 hfai.pl 插件来进行适配，具体包括：

hfai.pl.HFAIEnvironment，自动适配萤火集群的多卡并行环境，在训练时加入插件就可以正常使用；
hfreduce_bind_numa，使用 hfreduce 加速通信，绑定 buma 避免多卡间额外的网络开销；
hfai.pl.nn_to_hfai，使用 hfai 优化算子替换模型中的基础算子，加速训练。

具体操作如下：

在配置文件中将 trainer 的 strategy 指定为 hfreduce_bind_numa：

trainer:
    max_epochs: 300
    strategy: hfreduce_bind_numa
    ...

在训练代码中，使用 nn_to_hfai 算子加速和 HFAIEnvironment 的环境设置功能：

diffusionModelModule = nn_to_hfai(diffusionModelModule)

...

trainer = Trainer.from_argparse_args(trainer_opt, **trainer_kwargs)
trainer.plugins = [HFAIEnvironment()]

通过如上简单几步操作就可以将萤火集群的加速特性融入到 Stable Diffusion 的训练代码中。测试在使用 hfai.pl 前后模型训练速度的变化，可以发现模型的单次 forward 时长从 0.787 秒加快到了 0.758 秒，速度提升了 3.8%。

训练优化

我们在 Google Conceptual Caption 上使用 256x256 的分辨率进行训练，使用了在 ImageNet 上预训练的权重来初始化用来进行隐空间映射的 VQVAE 模型。

在训练时我我们尝试了使用 4，8，16，32 节点分别进行 Stable Diffusion 的训练以测试不同并行条件下的敏感度。在逐步扩大 Stable Diffusion 训练规模的过程中，我们发现 Stable Diffusion 对学习率非常敏感，然而学习率不能轻易随 Batchsize 增大而增大，这很容易导致梯度爆炸的发生。因此，我们在训练时采用了 Warmup 和 Gradient Clipping 的方法来帮助模型加速收敛，避免无法收敛的情况。

上图中橙色和蓝色曲线分别展示了有无使用 Warmup 和 Gradient Clipping 训练时的 loss 情况。如橙色曲线所示，当不使用 Clipping 和 Warmup 策略，模型在训练到第 400 个 step 的时候就出现了梯度爆炸的现象，无法继续正常进行收敛。而在使用之后，模型训练的学习率增长变得更加平缓，能够有效的避免发生梯度爆炸的发生。

训练结果

我们将模型在 Googlecc 数据集上训练了 240K Steps 后进行测试。在 COCO FID-30K（一个 COCO Caption 数据集中随机挑选的子集，由 3w 张图片组成）上，模型取得了 16.5 的 FID 指标，说明模型生成的图片能比较好的体现出文本中的内容。

以下是一些样例文本和将他们输入模型后生成的图片：

(a) A photo of a woman skiing on a white mountain.

(b) A painting of a squirrel eating a burger.

(d) A photo of a dog playing in a green field next to a lake.

可见虽然训练数据集规模较小，但模型仍然达到了较好的生成效果。

体验总结

Stable Diffusion 作为 AI 作图领域的旗舰模型，受到了广泛的关注，在小范围数据上训练也可以实现惊艳的生成效果。我们借助幻方萤火集群，通过简单几步改造，能比较轻松地实现 Stable Diffusion 的训练加速，证明了萤火集群的易用性和实力。综合体验打分如下：

研究指数：★★★★

作者提出了一种在隐空间上进行扩散的生成模型结构，降低了扩散模型运行开销的同时保证了生成质量。模型还应用了交叉注意力机制来辅助条件控制生成，并且支持多种不同模态条件下的图像生成。
开源指数：★★★★★

作为首个完全开源代码、训练数据和预训练权重的 AI 绘画预训练大模型，stable-diffusion 在学术界和其他相关领域都产生了极大的影响力。
门槛指数：★★

由于模型对资源占用有所优化，且开源工作完善，因此单个普通 GPU 即可运行模型推理。但训练开销较高。
通用指数：★★★★

作者提出的在隐空间上进行扩散的方法对一般的扩散模型都能够适用，并且基于交叉注意力的条件控制方法也能将模型应用于许多不同任务类型，对生成领域研究工作有广泛的借鉴意义。
适配指数：★★★

该项目依赖 pytorch-lightning，需要对萤火集群进行一定适配，但通过 hfai.pl 工具也能比较容易的在幻方 AI 环境运行并获得加速效果。

幻方 AI 紧跟 AI 研究的前沿浪潮，致力于用领先算力助力 AI 落地与价值创造，欢迎各方数据研究者与开发者们一同共建。

hfai 心法总章 | 快速上手攻略

High-Flyer — Fri, 16 Dec 2022 21:00:00 GMT

幻方 AI 发布了其沉淀多年的深度学习套件 hfai，吸引了众多同行研究员和开发者们咨询试用。整个套件的功能较多，而熟悉掌握了这套规则，是能够轻松地调用起平台的算力资源，从而高效完成训练任务的。为此，我们专门创建了“hfai 使用心法”系列专辑，陆续为大家介绍 hfai 一些功能的设计思路和原理，帮助大家更快上手 hfai，游刃有余得应对深度学习作业的各项挑战。

本终章为大家连贯的演示 hfai 心法在一个实际训练中的流畅应用。

官方文档：https://doc.hfai.high-flyer.cn/index.html

案例代码：https://github.com/HFAiLab/hfai-models

hfai workspace

打通本地与萤火之间的任督二脉。

其旨在同步本地工程目录代码到远程萤火超算中，更丝滑地利用萤火超算的算力运行您的项目工程。

hfai datasets

运转数据入萤火训练的不二法门。

一套专门为数据集量身定制的“神功”，解决大规模的训练数据，上传、管理和部署。

hfai python

任务提交任意所至，萤火训练行云流水。

其可以帮助大家方便快捷地发起、管理训练任务。

haienv

运行环境的影分身术。

其可以完成“本地 - 萤火，环境共享”——快速实现运行环境的集群部署。

hfai images

自建镜像，环境配置的最后绝招。

当面对一些特殊的场景，单单通过安装 Python 依赖已不足以满足需求，其可以帮助研究者构建项目特有镜像。

haiscale | 幻方萤火高性能并行训练工具库

小杞 — Tue, 06 Dec 2022 21:00:00 GMT

haiscale (Highflyer AI Scale) 是一个轻量级的高性能并行训练工具库，其整合了幻方 AI 多年的并行训练研发优化经验，能够帮助 PyTorch 用户更加高效、便捷地在大规模集群上训练模型。

haiscale 中包含了以下几种工具：

haiscale.ddp: 分布式数据并行工具，以幻方 AI 自研的 hfreduce 通信为后端，相比于 NCCL 能够获得更好的多卡拓展性能；
haiscale.fsdp: 极致优化 Fully Sharded Data Parallel (FSDP) 算法的实现，相比于 PyTorch FSDP 速度更快、占用显存更少；
haiscale.pipeline: 分布式流水线并行（或称模型并行）工具包，包含 GPipe, PipeDream 等算法，支持多机多卡训练；
haiscale.cpu_offload: 神经网络模型 Offload 工具，节省训练占用的显存。

下图展示了 haiscale 三种并行方式的性能，其相比 PyTorch 官方自带工具都有显著的性能提升：

用于测试的模型是 GPT-2 Medium，相关代码已开源至 hfai 模型仓库。下面将为大家简要介绍。

API文档：https://doc.hfai.high-flyer.cn/api/haiscale_ddp.html

示例模型：https://github.com/HFAiLab/hfai-models/tree/main/gpt

安装

haiscale 提供 Python 接口，通过如下方式安装：

如果要使用 haiscale DDP，首先需要先安装 hfreduce （如果不需要使用 DDP 可跳过这步）：

sudo apt install libnuma-dev
sudo apt install libibverbs-dev
pip install hfreduce --extra-index-url https://pypi.hfai.high-flyer.cn/simple --trusted-host pypi.hfai.high-flyer.cn

安装 haiscale：

pip install haiscale --extra-index-url https://pypi.hfai.high-flyer.cn/simple --trusted-host pypi.hfai.high-flyer.cn

haiscale.ddp

haiscale.ddp.DistributedDataParallel (haiscale DDP) 是一个分布式数据并行训练工具，使用 hfreduce 作为通讯后端，反向传播的同时会异步地对计算好的梯度做 allreduce。

haiscale DDP 的使用方式和 pytorch DDP 几乎相同，以下是使用示例：

from haiscale.ddp import DistributedDataParallel

model = MyModel().cuda()
model = DistributedDataParallel(model)
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)

# training ...
for step, (x, y) in enumerate(dataloader):
    optimizer.zero_grad()
    output = model(x)
    loss_fn(y, output).backward()
    optimizer.step()

# stop hfreduce
model.reducer.stop()

如果需要做梯度累加，可以使用 model.no_sync() 来减少通讯的开销。注意只需要最后一次反向传播时做 allreduce：

from haiscale.ddp import DistributedDataParallel

ddp = DistributedDataParallel(model, ...)
with ddp.no_sync():
    for input in inputs:
        ddp(input).backward()  # no synchronization, accumulate grads
ddp(another_input).backward()  # synchronize grads

haiscale.fsdp

Fully Sharded Data Parallel (FSDP) 是 META 在 ZERO-3 的基础上提出的分布式数据并行工具，它把模型的参数进行切分并分散到不同的 GPU 上，每块 GPU 上只有 1/ngpus 的参数。在做前向和反向传播时，FSDP 会先做 allgather 获得完整的参数，然后在前向和反向传播结束后释放掉，只保留 1/ngpus 的参数和梯度。FSDP 通过参数分片的方式，能够减少模型参数、梯度、优化器状态的显存占用，帮助我们训练更大规模的模型。

haiscale.fsdp.FullyShardedDataParallel 的使用方法和 DDP 类似，但优化器必须在 FSDP 之后创建，并且保存模型参数的时候需要先调用 summon_full_params。以下是使用示例：

from haiscale.fsdp import FullyShardedDataParallel

model = MyModel().cuda()
model = FullyShardedDataParallel(model)
optimizer = torch.optim.AdamW(model.parameters(), lr=0.01)

# training ...
for step, (x, y) in enumerate(dataloader):
    optimizer.zero_grad()
    output = model(x)
    loss_fn(y, output).backward()
    optimizer.step()

# save checkpoint
with model.summon_full_params():
    if rank == 0:
        state = model.state_dict()
        torch.save(state, 'model.pt')

haiscale FSDP 还支持传入 auto_wrap_policy 参数，具体作用可以参考 PyTorch FSDP 的文档以及我们提供的 GPT-2 示例。

haiscale.pipeline

如上图所示，haiscale.pipeline 工具包中提供了三种流水线并行的算法：

GPipe: 把模型切分成 ngpus 份，所有 microbatch 的前向传播结束之后再做反向传播；
PipeDream: 把模型切分成 ngpus 份，前向和反向传播交替执行（non-interleaved 1F1B）；
Interleaved1F1B: 把模型切分成 ngpus * num_model_chunks 份，前向和反向传播交替执行。

对于中等规模的模型（比如 GPT-2 Medium），我们推荐优先使用 PipeDream，相比于 GPipe 和 Interleaved1F1B 占用显存更少，速度更快。对于超大规模的模型（比如 GPT-3），我们可以更加均匀、细粒度的切分它，这时候推荐使用 Interleaved1F1B。

haiscale.pipeline 提供了一个统一的 forward_backward 接口，我们需要传入损失函数 criterion 和标签数据 labels，损失函数会通过 loss = criterion(*outputs, *labels) 的方式调用，forward_backward 接口会返回一个元组 (losses, outputs)，其中losses 代表每个 microbatch 的 loss 值，outputs 代表模型的输出。只有最后一个 rank 的进程能够获得 loss 和输出，其他进程得到的是 (None, None)。

下面通过示例展示 GPipe 和 PipeDream 的用法：

from haiscale.pipeline import GPipe, PipeDream, partition

dist.init_process_group(...)
torch.cuda.set_device(local_rank)
rank, world_size = dist.get_rank(), dist.get_world_size()
torch.manual_seed(12345)

def loss_fn(out, y):
    return ((out - y)**2).sum()

model = nn.Sequential(...)
model = partition(model, rank, world_size)

# chunks: number of microbatches
model = PipeDream(model.cuda(), chunks=32)
# or model = GPipe(model.cuda(), chunks=32)

for x, y in dataloader:
    losses, outputs = model.forward_backward(x, criterion=loss_fn, labels=(y,), return_outputs=True)
    if rank == world_size - 1:
        loss = losses.sum().item()  # losses: torch.Size([32])

# eval
with torch.no_grad():
    out = model(x)
    if rank == world_size - 1:
        # calculate metrics ...

以下是 Interleaved1F1B 的使用示例：

from haiscale.pipeline import Interleaved1F1B, partition

dist.init_process_group(...)
torch.cuda.set_device(local_rank)
rank, world_size = dist.get_rank(), dist.get_world_size()
torch.manual_seed(12345)

def loss_fn(out, y):
    return ((out - y)**2).sum()

model = nn.Sequential(...)
modules = partition(model, rank, world_size, num_model_chunks=2)  # len(modules) = 2
modules = [m.cuda() for m in modules]
model = Interleaved1F1B(modules, chunks=32)

for x, y in dataloader:
    losses, outputs = model.forward_backward(x, criterion=loss_fn, labels=(y,), return_outputs=True)
    if rank == world_size - 1:
        loss = losses.sum().item()  # losses: torch.Size([32])

# eval
with torch.no_grad():
    out = model(x)
    if rank == world_size - 1:
        # calculate metrics ...

数据并行和流水线并行组合

haiscale 还支持同时使用 DDP 和流水线并行。比如我们有 16 块 GPU，我们可以把这 16 块 GPU 划分成两个组，每个组有 8 块 GPU，然后两个组之间做数据并行，组内做流水线并行。

看如下使用示例：

from haiscale.ddp import DistributedDataParallel as DDP
from haiscale.pipeline import PipeDream, partition, make_subgroups

dist.init_process_group(...)
torch.cuda.set_device(local_rank)
rank, world_size = dist.get_rank(), dist.get_world_size()

dp_group, pp_group = make_subgroups(pp_size=8)

model = nn.Sequential(...)
model = partition(model, pp_group.rank(), pp_group.size())

model = DDP(model.cuda(), process_group=dp_group)
model = PipeDream(model, chunks=64, process_group=pp_group)

criterion = nn.MSELoss()

for x, y in dataloader:
    model.forward_backward(x, criterion=criterion, labels=(y,))

CPU Offload

除了以上的并行策略，在深度学习的训练过程中，我们常常会遇到显存不足的问题。

haiscale.cpu_offload.CPUOffload 能够帮助我们在训练中把一部分需要保存的中间变量移动到 CPU 内存上，然后在反向传播时把需要用到的 tensor 传输回 GPU 显存里，从而达到节省显存的目的。

haiscale 采用异步传输拷贝策略，能够把一部分的传输时间和 GPU 的计算重叠起来，从而减少拷贝带来的开销，提升整体计算效率。

使用时需要指定 offload_ratio 参数，其代表需要 offload 的中间变量的比例，offload_ratio=1 代表所有保存的中间变量都会被移动到 CPU 内存里。

以下是使用示例：

from haiscale.cpu_offload import CPUOffload

model = MyModel().cuda()
optimizer = torch.optim.AdamW(model.parameters(), lr=0.01)

for x, y in dataloader:
    optimizer.zero_grad()
    with CPUOffload(offload_ratio=0.1, tag="MyModel"):
        output = model(x)
    loss_fn(y, output).backward()
    optimizer.step()

萤火跑模型 | 多卡并行实现 YOLOv5 高性能训练

Vachel — Wed, 30 Nov 2022 21:00:00 GMT

YOLO 系列模型在整个深度学习目标检测领域有举足轻重的地位，近些年模型性不断发展，工程落地中的应用也十分广泛。

最近来自美国的 Ultralytics 公司发布了第五代 YOLO 模型: YOLOv5。相比上代模型，YOLOv5 以更轻量的参数，更极致的推理速度，一经发布受到了学术界和工业界相关人士的广泛关注。

幻方 AI 最近在萤火集群上对该项工作进行了体验和优化。YOLOv5 没有相应的论文说明，仅开源了代码。我们通过自研的 3FS、数据集仓库、hfreduce、算子等对模型训练进行优化，实现了近 50% 的训练提速。本文为大家简要描述。

项目地址：https://github.com/ultralytics/yolov5

模型仓库：https://github.com/HFAiLab/yolov5

模型简介

YOLOv5 是一种单阶段目标检测算法，下图展示了其整体结构框图：

对于一个目标检测算法而言，我们通常可以将其划分为4个通用的模块，具体包括：输入端、基准网络、Neck网络与Head输出端，对应于上图中的4个红色模块。YOLOv5 在 YOLOv4 的基础上添加了一些新的改进思路，使其速度与精度都得到了极大的性能提升。主要的改进思路如下所示：

输入端：表示输入的图片，该阶段通常包含一个图像预处理阶段，即将输入图像缩放到网络的输入大小，并进行归一化等操作。YOLOv5 提出了一些改进思路，如在网络训练阶段，使用 Mosaic 数据增强操作提升模型的训练速度和网络的精度；提出一种自适应锚框计算与自适应图片缩放方法；
基准网络：通常是一些性能优异的分类器种的网络，该模块用来提取一些通用的特征表示。融合其它检测算法中的一些新思路。YOLOv5 中不仅使用了 CSPDarknet53 结构，而且使用了 Focus 结构作为基准网络；
Neck网络：目标检测网络在 BackBone 与最后的 Head 输出层之间往往会插入一些层，利用它可以进一步提升特征的多样性及鲁棒性。虽然 YOLOv5 同样用到了 SPP 模块、FPN+PAN 模块，但是实现的细节有些不同；
Head输出层：用来完成目标检测结果的输出。针对不同的检测算法，输出端的分支个数不尽相同，通常包含一个分类分支和一个回归分支。YOLOv5 输出层的锚框机制与 YOLOv4 相同，主要改进的是训练时的损失函数 GIOU_Loss，以及预测框筛选的 DIOU_nms。

模型实践

幻方 AI 基于 Ultralytics 开源的 YOLOv5 代码，进行多级多卡的训练支持，开源至 HFAI YOLOv5，并采用幻方一系列优化工具进行提速升级，包括 hfai.datasets 高速数据集仓库，hfreduce 并行训练，hfai.nn 高性能算子等功能。

数据集

YOLOv5 开源代码里给出的训练案例，使用的训练数据只有 128 张图片构成的 COCO 子数据集。数据规模较小，且 Dataloader 读取训练样本时需要在存储系统里打开大量的小文件，整体训练吞吐会受到很大影响。

源代码里用 Yaml 配置训练数据集，通过如下方式构建 Dataloader:

from torch.utils.data import DataLoader, Dataset

class LoadImagesAndLabels(Dataset):
  def __init__(self, ...):

    f = []  # image files
    for p in path if isinstance(path, list) else [path]:
        p = Path(p)  # os-agnostic
        if p.is_dir():  # dir
            f += glob.glob(str(p / '**' / '*.*'), recursive=True)
        elif p.is_file():  # file
            with open(p) as t:
                t = t.read().strip().splitlines()
                parent = str(p.parent) + os.sep
                f += [x.replace('./', parent) if x.startswith('./') else x for x in t]  # local to global path
        else:
            raise FileNotFoundError(f'{prefix}{p} does not exist')
    self.im_files = sorted(x.replace('/', os.sep) for x in f if x.split('.')[-1].lower() in IMG_FORMATS)

    ...

可以看到，单个样本直接从本地存储系统中读取，如果使用 COCO 全量数据集，则需要打开的小文件数量对存储系统的压力不可小觑。

这里，我们采用自研的高速数据集仓库，hfai.datasets，优化训练数据的格式。将大量的小文件通过 FFRecord 整合转化为若干个集中的大文件，并且可以使用 FFRecord 中优化的随机读和批量读，进行数据样本的高速随机访问。同时，hfai 数据集仓库封装了很多常用的公开数据集，并且采用幻方自研的高速存储 3FS 进行保存，可以实现极大的性能优化。改造非常简单，如下例所示：

from hfai.datasets import COCODetection
from ffrecord import FileReader
from ffrecord.torch import Dataset, DataLoader

class LoadImagesAndLabels(Dataset):
  def __init__(self, ...):
    self.data_dir = COCODetection('train').data_dir
    self.reader = FileReader(self.data_dir / "train2017.ffr", check_data=True)

    ...

通过上面两行代码，就可以轻松加载 hfai 数据集仓库中的 COCODetection 数据集，大幅提升训练样本的读取效率。更多的数据集和使用，可以阅读 hfai.datasets 官方文档。

hfai 训练改造

萤火提供了统一的训练管理平台，可以将大量的训练任务依据优先级负载均衡到不同的GPU上执行，充分利用起计算算力。对于 YOLOv5 的训练，我们可以加大数据并行，极大加速训练速度。操作上只需要加入如下几行代码：

引入 hfai
```
import hfai
```

使用 hfreduce, hfai.nn 替换代码中的参数

import hfai.nn as hfnn
from hfai.nn.parallel import DistributedDataParallel

def main(local_rank, args):

  ...

  model = hfnn.to_hfai(model)
  model = DistributedDataParallel(model.cuda(), device_ids=[local_rank])

  ...

使用 hfai.multiprocessing.spawn 启动多卡并行训练，绑定 numa 避免网络拥塞以提高训练效率

if __name__ == "__main__":
  args = parse_args()
  ngpus = torch.cuda.device_count()
  hfai.multiprocessing.spawn(main, (args,), nprocs=ngpus, bind_numa=True)

训练每步都轮询一次萤火集群，看当前的计算资源是否足够继续训练，做好断点的保存

start_epoch, start_step, model, optimizer, ... = load_ckpt()

def train(...):
  for epoch in range(start_epoch, epochs):
    for step, batch in enumerate(dataloader):
      if step < start_step:
        continue
      
      ...

      if rank == 0 and local_rank ==0 and hfai.client.receive_suspend_command():
        save_ckpt(model, epoch, step+1, optimizer, ...)
        hfai.go_suspend()

通过如上4步简单的改造，我们就可以基本利用起萤火集群的特性，对 YOLOv5 的训练进行优化提升。更多详细的代码改造细节，可以阅读我们在模型仓库中开源的代码。

训练结果

我们申请 4 个萤火计算节点进行并行训练。通过 hfai 命令行工具提交任务，测试萤火集群对 YOLOv5 训练加速的性能。如下图所示：

可以看到，左右两图分别展示了萤火集群特性（hfai）相比原始训练方式（raw）和不使用算子，hfreduce (hfai w/o nn&hfreduce) 的加速效果。如左图所示，完整的萤火集群特性相比原始训练方式有将近 50% 的训练提速，这得益于优秀的训练数据集管理和优良的算子与通信能力；如右图所示，在使用同样的数据加载模式下，hfreduce 和 hfai 优化算子可以进一步提升训练速度，充分利用显卡计算资源，整体获得将近 20% 的训练提速。

下面给大家展示一些萤火集群训练出的 YOLOv5 模型在目标检测任务上的效果：

体验总结

YOLOv5 作为目标检测领域的旗舰模型，受到了广泛的关注。其不提供论文直接开源代码，并且用 Yaml 配置文件来指导深度学习模型的构建，在深度学习领域中也是独树一帜。我们借助幻方萤火集群，很轻松地实现了集群特性与这套新的开发模式的打通，获得明显的加速效果，证明了萤火集群的易用性和实力。

综合体验打分如下：

研究指数：★★★

该模型是目标检测领域下的旗舰模型，成熟且已在工业界成功地落地实践。
开源指数：★★★★★

数据和代码都已开源，代码逻辑清晰可读性高。
门槛指数：★★★

样例数据量小，模型规模不大。代码开发模式新颖。
通用指数：★★

模型只针对目标检测场景，其他场景待验证。
适配指数：★★★★★

依赖简单，PyTorch 框架构建的模型，只需要修改几行代码就能在萤火集群上执行。

幻方 AI 紧跟 AI 研究的前沿浪潮，致力于用领先算力助力AI落地与价值创造，欢迎各方数据研究者与开发者们一同共建。

hfai.pl | 兼具萤火集群的优化特性的 PyTorch Lightning

Freja — Thu, 17 Nov 2022 21:00:00 GMT

Pytorch Lightning（简称 pl）是在 PyTorch 基础上进行封装的库，它能帮助开发者脱离 PyTorch 一些繁琐的细节，专注于核心代码的构建，在 PyTorch 社区中备受欢迎。本文为大家介绍如何在 PyTorch Lightning 上轻松的应用幻方萤火集群的系列优化特性。

集群环境适配

若集群每个计算节点有 x 张 GPU。用户提交任务时需选定节点数量 N，则该任务可获得 N*x 个 GPU。每个进程中全局环境变量含义如下：

world_size (hfai): 节点数量，用 N 表示
rank (hfai): 节点 id，用 n 表示，n 属于 0 ~ N-1
local_rank (hfai): GPU id, 用 k 表示，k 属于 0 ~ x-1

与 PyTorch init_process_group 的变量有如下对应：

world_size (PyTorch): 进程数量，每个 GPU 对应一个进程，因此总进程数目为总 GPU 数目，计算方法为 N*x
rank (PyTorch): 进程 id，利用节点 id 和 GPU id 可以计算出进程的 id，计算方法为 n*x+k

在一般的 PyTorch 代码中，我们需要进行如下的分布式初始化：

ip = os.environ.get("MASTER_ADDR", "127.0.0.1")
port = os.environ.get("MASTER_PORT", "2223")
hosts = int(os.environ.get("WORLD_SIZE", 1))  # number of nodes
rank = int(os.environ.get("RANK", 0))  # node id
gpus = torch.cuda.device_count()  # gpus per node
dist.init_process_group(
    backend="nccl", 
    init_method=f"tcp://{ip}:{port}", 
    world_size=hosts*gpus, 
    rank=rank*gpus+local_rank
)
torch.cuda.set_device(local_rank)

pl 对分布式初始化进行了封装，用户不需要自己初始化，只需指定节点数量和 GPU 数量即可。为了适配到幻方萤火集群，我们提供了 HFAIEnvironment 环境插件，只需要几行代码即可将 pl 的分布式参数无缝适配到萤火集群当中。

from hfai.pl import HFAIEnvironment
trainer = pytorch_lightning.Trainer(
    gpus=x, num_nodes=N,
    plugins=[HFAIEnvironment()] # 定义 HFai 环境并作为插件输入
)

绑定 numa

为了规避跨 numa 传输带来的带宽损耗，建议对训练代码进行 numa 绑定，在一般的 PyTorch 代码中，我们通过使用 hfai.multiprocessing 启动多进程，并指定 bind_numa=True 来进行 numa 的绑定，代码如下：

import torch
import hfai

def main(gpu_id):
    torch.cuda.set_device(gpu_id)
    # ......

if __name__ == "__main__":
    ngpus = hfai.utils.num_gpus()  # 调用 cuda 函数会导致子进程产生错误
    hfai.multiprocessing.fork(main, args=(), nprocs=ngpus, bind_numa=True)

pl 中对多进程的启动进行了封装，用户只需要在 strategy 中指定使用 ddp 或者 ddp_spawn 就可以指定多进程的启动方式。因此我们提供了两种新的 strategy：ddp_bind_numa 和 ddp_spawn_bind_numa，用户选择这两种 strategy，就可以在指定多进程启动方式的同时绑定 numa，使用方式如下所示：

# 使用 ddp_bind_numa 或者 ddp_spawn_bind_numa
trainer = pytorch_lightning.Trainer(strategy="ddp_spawn_bind_numa")

使用 hfreduce

hfreduce 是幻方 AI 自研的高性能多卡并行通信工具，其能够更高效的在多显卡之间交换梯度信息，加速模型训练。

在一般的 PyTorch 代码中，我们通过使用 hfai.nn.parallel.DistributedDataParallel 替换 PyTorch 自带的 torch.nn.parallel.DistributedDataParallel 即可使用 hfreduce。

由于 pl 中对 DDP 的初始化进行了封装，因此我们提供了 hfreduce_bind_numa 和 hfreduce_spawn_bind_numa 两种 strategy。用户在初始化 Trainer 的时候进行如下指定即可。

# 使用 hfreduce_bind_numa 或者 hfreduce_spawn_bind_numa
trainer = pl.Trainer(strategy="hfreduce_bind_numa")

自动打断重启

幻方萤火平台采用任务级分时调度的底层设计，给每个任务分配集群运行时间片。按照分时调度的方案规则，训练任务会被中断、并自动调起。因此在该平台上运行的代码需要在训练循环中进行打断信号的捕获和参数保存，以便之后恢复训练。断点保存的示例代码如下：

for epoch in range(epochs):
    for step in range(len(data_batch)):
        if hfai.distributed.get_rank() == 0 and gpu_id == 0: # 在0号节点的0号进程上接收集群调度信息
            if hfai.client.receive_suspend_command(): 
                model.save() # 保存模型、迭代器等参数到文件
                time.sleep(5) # 最多预留5秒完成断点保存，之后会被强制打断
                hfai.client.go_suspend() # 发送准备好被打断的信号

pl 对训练循环进行了封装，因此我们提供了 ModelCheckpointHF 作为回调，在收到打断信号的 step 进行模型的保存，在每次任务启动时，检测是否需要从上一个断点恢复训练，使用方法如下所示：

from hfai.pl import ModelCheckpointHF
output_dir = 'hfai_out'
cb = ModelCheckpointHF(dirpath=output_dir)
trainer = pytorch_lightning.Trainer(callbacks=[cb]) # 自动处理集群打断信号

ckpt_path = f'{output_dir}/{cb.CHECKPOINT_NAME_SUSPEND}.ckpt' # 检查是否有断点模型被保存
ckpt_path = ckpt_path if os.path.exists(ckpt_path) else None
trainer.fit(
    model_module,
    ckpt_path=hfai_suspend_ckpt_path # 自动恢复训练
)

使用 hfai 优化算子

幻方 AI 依托萤火二号集群，对 PyTorch 框架进行了深度优化，结合萤火集群的特点，对一些常用的 AI 算子重新研发，提升效率，进一步提升了模型整体的训练效率。

在 pl 框架中，我们只需要增加如下一行代码，就可以将 PyTorch 中的算子转换成 hfai 优化后的算子：

import pytorch_lightning as pl
class ToyNetModule(pl.LightningModule):
    # ...
pl_module = ToyNetModule()
model_module = nn_to_hfai(pl_module) # 将算子转换为 hfai 算子

使用 FFRecord

为了在 PyTorch Lightning 中使用 ffrecord 的 Dataloader，我们需要在 Dataloader 设置 skippable=False:

from ffrecord.torch import Dataset, DataLoader

class MyDataset(Dataset)
    # ...

dataset = MyDataset(...)
dataloader = DataLoader(dataset, batch_size, num_workers=num_workers, skippable=False)

使用示例

上面介绍了如何将 pl 框架的代码融入萤火集群的各种优化特性当中，下面提供一个完整的示例帮助大家理解：

from hfai.pl import HFAIEnvironment
from hfai.pl import ModelCheckpointHF
import pytorch_lightning as pl

class ToyNetModule(pl.LightningModule):
    ...

output_dir = 'hfai_out' # 模型保存路径
cb = ModelCheckpointHF(dirpath=output_dir) # 可以接收集群打断信号的回调类
trainer = pl.Trainer(
    gpus=x, # 每个节点 x 个 GPU
    num_nodes=N, # 节点数量
    strategy="ddp_bind_numa",  # 支持 ddp_bind_numa, ddp_spawn_bind_numa, hfreduce_bind_numa, hfreduce_spawn_bind_numa
    plugins=[HFAIEnvironment()],  # 自动适配 HFAI 分布式环境
    callbacks=[cb] # 自动处理集群打断信号
)
model_module = nn_to_hfai(ToyNetModule()) # 将算子转换为 hfai 算子

ckpt_path = f'{output_dir}/{cb.CHECKPOINT_NAME_SUSPEND}.ckpt' # 判断之前是否保存了断点模型
ckpt_path = ckpt_path if os.path.exists(ckpt_path) else None
trainer.fit(
    model_module,
    ckpt_path=ckpt_path # 自动恢复训练
)

通过上面的简单适配，PyTorch Lightning 就能够应用幻方萤火集群的系列优化特性啦。

hfai images | 自定义镜像，解决环境配置的最后一里路

High-Flyer — Fri, 11 Nov 2022 21:00:00 GMT

幻方 AI 发布了其沉淀多年的深度学习套件 hfai ，吸引了众多同行研究员和开发者们咨询试用。整个套件的功能较多，而熟悉掌握了这套规则，是能够轻松地调用起平台的算力资源，从而高效完成训练任务的。为此，我们专门创建了 “hfai 使用心法” 系列专辑，陆续为大家介绍 hfai 一些功能的设计思路和原理，帮助大家更快上手 hfai，游刃有余的应对深度学习作业的各项挑战。

之前的文章为大家介绍了使用 haienv 构建运行环境和执行任务，然而对于一些特殊的场景，如强化学习、物理仿真等，单单通过安装 Python 依赖已不足以满足需求。研究者往往会构建项目特有镜像。如何将这类需求在萤火集群上支持是幻方 AI 研发者们在思考研究的问题。

最近，hfai images 功能面世，为研究者和开发者们提供了一套解决此类方法的绝招。本期文章将为大家详细介绍。

使用场景

如上一期文章所介绍的，haienv 为萤火集群的用户提供了自建 Python 运行环境的能力，其支持 pip, conda 两种方式满足大部分用户环境构建的需求。

而 hfai images 解决的是一些非 Python 环境安装的问题，比如物理仿真里要安装模拟器，分子计算工具等。

幻方 AI 对集群中的镜像进行了深度优化和封装，以满足萤火集群高性能的深度学习训练要求。因此 hfai 虽支持用户的自有镜像构建，但需要符合 hfai 的规范和流程，其过程需要幻方管理员的参与和支持。

在确定使用 hfai images 之前，请您审视自己的环境安装需求，确定只有自有镜像这一条路比较方便后，可以和管理员联系。

使用流程

hfai images 的使用流程包含如下几步：

获取集群中的 base 镜像的 tar 包。我们提供了 cuda_111 和 cuda_113 两个版本，基于 Ubuntu 20.04；
在本地加载上述 tar 包，比如 registry.high-flyer.cn/platform/hf_training:ubuntu2004-cu113-20221019；
基于上述 base 镜像构建自己的镜像（比如 python:3.8-alpine）；
docker save 自己的镜像，保存成 tar 包（比如 python-alpine.tar），通过 hfai workspace 上传到萤火集群；
打开 Studio 进入 workspace 目录，使用 hfai images load 命令（比如 hfai images load python-alpine.tar）；
等待完成，可以通过 hfai images list 命令查看。这里镜像加载完成后返回如下结果：
等状态变成 loaded 之后，您才可以通过 hfai python test.py -- -n 1 --image registry.high-flyer.cn/hfai/python:3.8-alpine 来使用您所构建的镜像。

在 MM 系列框架上使用 FFRecord

小杞 — Thu, 03 Nov 2022 21:00:00 GMT

FFRecord 是幻方AI提供的训练数据格式，其可以整合大量零散的小文件，支持训练数据的随机索引，极大发挥分布式高速存储的性能，提高训练效率。然而，对于 mmdetection, mmsegmentation 等广受研究者欢迎的 mm-* 系列框架，其深度封装了数据和模型，让萤火集群的用户不能方便快捷的进行 FFRecord 转化。对此，幻方AI针对这类场景进行了优化，提供了 FFRecordClient 接口。用户能够在配置文件中选择使用 FFRecordClient 来替换默认的读取后端。

那么 FFRecordClient 和之前的 FFRecord 有什么不同？它具体该如何使用？本期文章将为大家分享。在阅读本文之前，建议读者先了解一下 MMCV FileClient 的原理。

方法介绍

我们知道，FFRecord 可以将大量的小文件整合到几个大文件中。与 TFRecord 这种开源文件格式不同，FFRecord 可以支持随机读取，批量读取，可以极大地提高高速存储的数据读取性能。对于 mm-* 系列框架中的数据封装，FFRecord 扩展了文件夹打包能力，以适配这种场景。具体的使用方式如下：

import ffrecord
ffrecord.pack_folder("/path/to/dataset", "packed.ffr")

在训练中，可以使用 FFRecordClient 读取后端

在训练代码中导入 FFRecordClient：
```
import hfai.utils.mm
```

修改配置文件，为所有的读取小文件的操作添加 file_client_args 参数，比如在 mmseg 中：

file_client_args=dict(
    backend="ffrecord",
    fname="packed.ffr",
)

train_pipeline = [
    dict(type='LoadImageFromFile', file_client_args=file_client_args),
    dict(type='LoadAnnotations', reduce_zero_label=True, file_client_args=file_client_args),
    dict(type='Resize', img_scale=(2048, 512), ratio_range=(0.5, 2.0)),
    ...
]

使用示例

接下来我们展示两个示例，分别是在 mmdetection 和 mmsegmentation 上的使用。

mmdetection

以 coco 数据集为例，假设我们已经有 coco 的原始数据集，目录结构如下：

coco/
├── annotations
├── train2017
└── val2017

我们先打包整个数据集到 data/coco/coco.ffr 中：

import ffrecord
ffrecord.pack_folder("coco/", "data/coco/coco.ffr")

然后把 annotations 文件夹单独拷贝出来放到 data/coco/annotations，现在目录结构如下：
```
data
└── coco
    ├── annotations
    └── coco.ffr
```
在训练代码中导入 FFRecordClient：
```
import hfai.utils.mm
```
修改配置文件 configs/_base_/datasets/coco_detection.py 的内容为： coco_detection.py

mmsegmentation

以 ADE20k 数据集为例，假设我们已经有 ADE20k 的原始数据集，目录结构如下：

ADEChallengeData2016
├── annotations
├── images
├── objectInfo150.txt
└── sceneCategories.txt

我们先打包整个数据集到 data/ade20k.ffr 中：

import ffrecord
ffrecord.pack_folder("ADEChallengeData2016/", "data/ade20k.ffr")

在训练代码中导入 FFRecordClient：
```
import hfai.utils.mm
```
修改配置文件 configs/_base_/datasets/ade20k.py 的内容为： ade20k.py

萤火跑模型 | 能量函数指导的图图翻译扩散模型

Chongxuan Li — Thu, 20 Oct 2022 21:00:00 GMT

扩散模型作为一种新的深度生成模型，在图像生成领域取得了 SOTA 的效果，并且逐渐在诸多应用领域展现出强大的性能，如视频生成、分子图建模等。

最近来自人民大学李崇轩课题组的研究者们提出了一种基于能量函数的扩散模型生成指导方法 EGSDE。该方法旨在通过预定义的能量函数对预训练好的扩散模型进行指导，从而实现可控生成。现有指导生成方法如 classifier-guidance 等可以理解为 EGSDE 的特例。理论上，研究者从专家产物的角度解释了 EGSDE 的采样分布。

该项目利用了幻方AI深度学习训练平台的算力及加速性能，在无监督的图到图翻译任务上进行了大量实验，在若干数据集上取得了最好的“可信-真实”权衡结果。该项工作被人工智能顶级学术会议 NeurIPS 2022 收录，本文将为大家简要描述。

论文标题：EGSDE: Unpaired Image-to-Image Translation via Energy-Guided Stochastic Differential Equations

论文地址：https://arxiv.org/abs/2207.06635

模型仓库：https://github.com/ML-GSAI/EGSDE

模型介绍

1. 研究动机

无监督的图到图翻译的任务是在没有成对数据的情况下学习源域到目标域的映射，且希望翻译后的图像一方面要去掉原图的域特有特征使其更“真实”，另一方面可以保留原图的域无关特征使其更“可信”。例如将猫翻译到狗，要把原图的鼻子、胡须这种特有特征去掉，同时保留原图的姿势、色调这种无关特征。

随着扩散模型的发展，逐渐开始涌现一些基于扩散模型的图到图翻译方法并且取得了 SOTA 的结果。这些方法往往是在目标域训练一个扩散模型保证“真实”，在测试的阶段用上原图的信息实现“可信”，但是并没有用上源域的训练数据。

基于此，该工作的研究者们希望同时利用源域和目标域的数据训练一个能量函数，通过能量函数对预训练好的目标域扩散模型进行指导生成，从而实现更“可信”且“真实”的图到图翻译。整体结构如下图所示：

2. 模型方法

该方法的核心要义是如何根据任务设计合适的能量函数。在图像翻译任务中，研究者将能量函数 $\varepsilon(y,x,t)$ 拆成两项对数势函数，分别负责生成图像的“真实”和“可信”。

\begin{aligned} d_y = [f(y,t)-g(t)^2(s(y, t) - \nabla_y \varepsilon(y,x_0,t))]dt + g(t)d\bar{w} \end{aligned}

\begin{aligned} \varepsilon(y,x,t) &= \lambda_s\varepsilon_s(y, x, t) + \lambda_i\varepsilon_i(y, x, t) \\ &= \lambda_s\mathbb{E}_{q_{t | 0}(x_t | x)}\varepsilon_s(y, x_t, t) - \lambda_i\mathbb{E}_{q_{t | 0}(x_t | x)}\varepsilon_i(y, x_t, t) \end{aligned}

进一步，研究者将负责“真实”的对数势函数 $\varepsilon_s(y,x,t)$ 定义为原图和生成的图片经过域独特的特征提取器 $\mathbb{E}_s(\cdot,\cdot)$ 后之间的相似度（余弦相似度），负责“可信”的 $\varepsilon_i(y,x,t)$ 则定义为两者经过域无关的特征提取器 $\mathbb{E}_i(\cdot,\cdot)$ 后之间的负相似度（负 $L_2$ 距离的平方）。

\begin{aligned} \varepsilon_s(y, x_t, t) &=cos(\mathbb{E}_s(y,t), \mathbb{E}_s(x_t, t)) \\ \varepsilon_i(y, x_t, t) &=-||\mathbb{E}_i(y,t), \mathbb{E}_i(x_t, t)||_2^2 \end{aligned}

直观上来说，去噪过程中能量函数的减少一方面导致原图和翻译图片间的特有特征相似度变低，从而去掉原图中源域特有的特征，另一方面导致两者的域无关的特征相似度变高，从而保留原图中域无关的特征。

实现上，如上图所示，研究者用分类器除去最后一层的剩下部分和低通滤波器分别当做域独特的特征提取器 $\mathbb{E}_s(\cdot,\cdot)$ 和域无关的特征提取器 $\mathbb{E}_i(\cdot,\cdot)$ 。

至此，能量函数和扩散模型已知，从一个带噪声的原图出发，可以通过扩散模型对 EGSDE 进行采样实现图到图的翻译。除此之外，将 $\mathbb{E}_s(\cdot,\cdot)$ 依赖的二分类器拓展到多分类器，该方法也很容易拓展到多域翻译任务。

理论上，研究者展示了对 EGSDE 的离散化采样近似等价于从如下专家产物形式的分布采样:

\begin{aligned} \tilde{p}(y_t|x_0) = \frac{p_{r1}(y_t|x_0)p_{r2}(y_t|x_0)p_{f}(y_t|x_0)}{Z_t} \end{aligned}

这为 EGSDE 中每个部分的作用提供了解释，其中扩散模型和 $\varepsilon_s(y,x,t)$ 对应于负责“真实”的专家， $\varepsilon_i(y,x,t)$ 对应于负责“可信”的专家。

实验

在实验中，研究者选择 FID 作为对“真实”的衡量，L2、PSNR、SSIM 作为对“可信”的衡量，人工测量 AMT 作为同时两方面的衡量，在 AFHQ 和 CelebA-HQ 数据集验证方法的有效性。如下表所示，EGSDE 取得了最好的“可信-真实”权衡结果，且通过调整超参可以进一步改善 FID。

为了验证 DAT 的有效性，研究者们展示了一些可视化结果。

如上图所示，研究者在 “Male→Female” 任务上展示了关于能量函数权重的消融实验， $\lambda_s$ 越大，生成的图片越像女性、越“真实”， $\lambda_i$ 越大，生成的图片保留原图更多的信息、越“可信”，验证了能量函数的作用。

总结

研究者们提出了一种基于能量函数的扩散模型生成指导方法 EGSDE。该方法旨在通过预定义的能量函数对预训练好的扩散模型进行指导，从而实现可控生成。该方法实现了更“可信”且“真实”的图到图翻译，同时可以推广应用到更多的场景领域。

同时，研究者采用幻方 AI 的一系列深度学习自研套件，极大加速了模型的训练过程，获得了大量的实验结果。我们欢迎更多优秀的课题与幻方合作，一道推动 AI 技术的发展与落地。

综合体验打分如下：

研究指数：★★★★

该模型改进提升了扩散模型方法，推动了图图翻译场景下的技术发展。
开源指数：★★★★★

数据和代码都进行了完整且规整的开源。
门槛指数：★★★

数据规模，模型大小适中，适合多级多卡数据并行训练。一般单卡训练难度比较大。
通用指数：★★★★

该方法能适用于很多可控生成任务，如材料和药物领域新分子的生成。
适配指数：★★★★★

依赖简单，很容易与幻方 AI 的训练优化工具结合，提效明显。

幻方 AI 紧跟 AI 研究的前沿浪潮，致力于用领先算力助力 AI 落地与价值创造，欢迎各方数据研究者与开发者们一同共建。

萤火跑模型 | 英伟达元宇宙创造营 GET3D

Freja — Tue, 18 Oct 2022 21:00:00 GMT

随着多个行业朝着大规模的 3D 虚拟世界发展，能够生成大量的、高质量的、多样的 3D 内容的工具是非常被需要的。英伟达的最新工作 GET3D 希望训练更好的 3D 生成模型，来生成下游任务可以直接使用的、保真纹理和复杂几何细节的 3D 模型。

幻方 AI 最近对这项工作进行了整理和优化，在幻方萤火二号上复现了实验。通过幻方自研的 3FS、hfreduce、算子，对模型训练进行提速，从单机多卡的训练支持优化到多级多卡的训练支持，帮助研究者和开发者们降低使用门槛。本期文章将为大家详细描述。

论文标题：GET3D: A Generative Model of High Quality 3D Textured Shapes Learned from Images

原文地址：https://arxiv.org/pdf/2209.11163.pdf

项目主页：https://nv-tlabs.github.io/GET3D

模型仓库：https://github.com/HFAiLab/GET3D

模型介绍

3D 生成模型

最近，英伟达发布了最新的 GET3D 模型，可以通过 2D 图像训练，生成具有高保真纹理和复杂几何细节的 3D 形状。在此之前，学术界也有大量的 3D 生成模型，他们所生成的 3D 模型存在如下问题：

缺乏几何细节
缺乏纹理
在合成过程中只能使用神经渲染器，3D 软件中不方便使用

过去工作的方法以及生成的 3D 模型效果如下表所示，其中 Ours 代表 GET3D 模型的生成几何结果，Ours-tex 代表 GET3D 模型生成的带纹理结果：

GET3D

GET3D 希望克服上述模型存在的问题，生成具有丰富几何细节、带纹理的、下游 3D 软件中可以直接使用的 3D 模型。并且训练过程只需要 2D 的监督图像训练。

GET3D 包括生成器部分和辨别器部分，生成网路包括如下两个分支：

几何生成分支：可微的输出任意拓扑的表面几何结果
纹理生成分支：根据查询的表面点来产生纹理场（texture field），还可以扩展到表面的其他属性，比如材质

生成器生成 3D 模型及纹理后，通过一个有效的可微栅格器，将生成的带纹理 3D 模型投影到 2D 的高分辨率图片，并使用 2D 辨别器来区分生成器所生成模型的投影图像和真实图像、生成器所生成模型的投影轮廓和真实轮廓。整个过程都是可微分的，使得整个对抗训练可以从辨别器传递到两个分支。模型整体结构如下图所示。

GET3D 输入是两个采样的高斯分布（z1 和 z2），通过非线性映射网络得到两个隐含表示（w1 和 w2），非线性网络是 8 层的 MLP 网络，每层是 512 维和 leaky-ReLU 的激活函数。两个隐含表示的具体含义如下：

纹理隐含向量（w1）：用来控制 3D 模型的形状
几何隐含向量（w2）：用来控制 3D 模型的纹理

3D 几何生成器：GET3D 的几何生成器包含最近提出的可微分表面表征 DMTet，DMTet 将 3D 模型表示成一个可变形三角面片的四面体（tetrahedron）及每个顶点的符号距离场（signed distance field，SDF），通过移动顶点来使得几何形状变形，达到从四面体可微分的恢复 3D 模型表面的目的。通过使用 DMTet，生成器可以生成任意拓扑结构和类别的 3D 模型。

纹理生成器：GET3D 将纹理参数化为纹理场（texture field），纹理生成器以纹理隐含向量（w2）和几何隐含向量（w1）为条件，将具体的表面点，映射到 RGB 表示的颜色空间中。

可微分渲染和训练：GET3D 将生成的 3D 几何模型和纹理场，使用可微分渲染器渲染到 2D 的图片，再使用 2D 辨别器来监督辨别器的学习。辨别器的学习目标是更准确的分辨真实物体的图片和生成物体的渲染图片。

模型实践

幻方 AI 基于英伟达开源的 GET3D 代码，进行多级多卡的训练支持，开源至 HFAI GET3D ，并采用幻方一系列优化工具进行提速升级，包括 hfreduce 并行训练、hfai 集群训练断点挂起等功能。

数据集

GET3D 使用 ShapeNET 数据集进行试验，ShapeNet数据集是一个有丰富标注的、大规模的3D图像数据集。ShapeNetCore 是 ShapeNet 的一个子集，其中包括近 51,300 个独特的 3D 模型。它提供了 55 个常见的对象类别和注释。GET3D 使用 ShapeNetCoreV1 数据集作为训练及测试数据。

我们对 ShapeNET 数据进行了渲染，将 3D 模型全部投影为 2D 图片以备训练。

模型训练

萤火提供了统一的训练管理平台，可以将大量的训练任务依据优先级负载均衡到不同的GPU上执行，充分利用起计算算力。我们只需加入如下几行代码，便可以把训练任务提交给萤火集群运行。

登录幻方萤火二号，引入 hfai
```
import hfai
```

初始化 hfreduce 分布式参数

from hfai.nn.parallel import DistributedDataParallel
model = DistributedDataParallel(model, device_ids=[local_rank])
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)

# training ...
for epoch in range(epochs):
    for step, (x, y) in enumerate(dataloader):
        # training
        optimizer.zero_grad()
        output = model(x)
        loss_fn(y, output).backward()
        optimizer.step()

对于每一轮训练，加入接收集群调度的逻辑代码，并做好模型 checkpoint 的保存

start_epoch, start_step, others = hfai.checkpoint.init(model, optimizer, ckpt_path='latest.pt')
for epoch in range(start_epoch, epochs):
    for step, (x, y) in enumerate(dataloader):
        if step < start_step:
            continue

        output = model(x)
        loss_fn(y, output).backward()
        model.try_save(epoch, step, others=None)

训练结果

我们申请 4 个节点，32 张 A100 显卡进行并行训练。通过hfai命令行工具提交任务，在萤火二号上，以车类别作为数据来训练 GET3D，每个 64 个 kimg 仅耗时 10 分钟左右，相比单卡性能提升了 32 倍以上。GPU 利用率在 95% 左右，充分利用其了显卡的计算性能。模型在训练到第 5000 个 kimg 基本达到了收敛状态。

训练效果可视化如下：

体验总结

GET3D 模型通过 2D 图像训练，生成具有高保真纹理和复杂几何细节的 3D 形状，解决了以往生成模型几何细节缺失、纹理缺失或者下游引擎无法使用等问题。能够生成大量的、高质量的、多样的 3D 内容，帮助多个行业更快速的朝着大规模的 3D 虚拟世界发展。幻方 AI 对 GET3D 代码适配到了多级多卡的训练，借助萤火平台对 GET3D 的训练进行了优化，最终达到了超过单卡 32 倍的训练速度。

综合体验打分如下：

研究指数：★★★★

该模型以低成本、2D 监督、端到端的方法实时生成 3D 模型，生成模型几何细节丰富、纹理自然。
开源指数：★★★★

数据处理和代码都已经开源，代码逻辑清晰、可读性高，HFAI 对其增加了多级多卡的训练支持。
门槛指数：★★

数据量、模型大小中等，普通高性能单卡即可运行，但是正常训练需要花费一周左右。
通用指数：★★★★

该方法使用 2D 领域生成网络的思想进行 3D 模型的生成，未来具有很大的拓展空间。
适配指数：★★★★★

该项目依赖简单，很容易与幻方AI的训练优化工具结合，提效明显。

幻方 AI 紧跟 AI 研究的前沿浪潮，致力于用领先算力助力AI落地与价值创造，欢迎各方数据研究者与开发者们一同共建。

萤火跑模型 | 分布式训练大规模深度图神经网络

Vachel — Tue, 11 Oct 2022 15:00:00 GMT

图（Graph）数据在现实世界中非常常见，例如社交网络、交通网络、物理系统等等，近几年图神经网络的发展将图数据的分析与深度神经网络结合，在越来越多的领域发挥出重要的作用，例如电商推荐、生物化学结构分析、反恐反诈风险控制等等。数据规模也呈现越来越大之势，动辄上千万节点规模的图，让很多图神经网络的训练需要探寻并行计算的方式以加速。

然而，图数据与视觉、自然语言等领域的数据不同，没有划分好的训练样本，很多情况下训练样本是从一个巨大的图中去随机采样，特征也需要在采样的过程中遍历邻接节点，这种天然特性使得对大规模图的并行训练并不容易。

幻方AI最近尝试对图神经网络的并行训练进行优化，以 OGB 图数据集作为训练样本，在幻方萤火集群上复现 DeepGCNs 的实验。基于幻方自研的3FS、hfreduce等工具，采用灵活的分布式训练方法，进行不同任务场景下的 GNN 训练提速，取得预期效果。本期文章将为大家详细描述。

论文标题：DeepGCNs: Can GCNs Go as Deep as CNNs?

原文地址：https://www.deepgcns.org/

模型仓库：https://github.com/HFAiLab/Distributed-DeepGCNs

数据集仓库: https://doc.hfai.high-flyer.cn/api/datasets.html#hfai.datasets.OGB

模型介绍

DeepGCNs 集成了近两年一系列前沿的 GNN 工作，其旨在构建如 CNN 一样的深层图卷积网络，解决优化深层结构中的一系列问题，以达到更优的表现效果。具体包括：

尝试构建深层的图卷积网络（GCN），探索深层 GCN 结构导致的信息传递快，容易梯度消息、过渡平滑和过渡拟合的问题；
设计提出残差、密集连接和空洞卷积等方案，使得深层图卷积网络表示更加稳定；
设计提出使用新的聚合、归一化函数，以统一信息的传递；
研究逆连接、组卷积、权重绑定和平衡模型，以提高深层图卷积模型的内存和参数效率。

DeepGCNs 的结构如下图所示：

在 OGB 公开图数据集的很多任务测试中，DeepGCNs 都取得了优异的效果。

可以看到，大规模图与深度图卷积模型的结合在越来越多的场景中发挥出惊人的效果，在 OGB 的公开结果测试排名中可以看到，对显卡算力的需求越来越大。如何高效的利用计算资源，尽可能加速图卷积网络的训练，是之后走向落地商用的必研究课题。

图的分布式训练

那么，如何对图数据进行分布式训练呢？在这之前，我们先看下图领域下的三大任务：

点属性预测（Node Property Prediction）
边属性预测（Link Property Prediction），也叫作链接预测
图属性预测（Graph Property Prediction）

这里图属性预测任务下，训练样本数据的构建方式和一般视觉、自然语言任务类似，输入的单个样本是一个完整的图，分布式训练可以在样本层面做拆分。具体方式如下例所示：

from torch.utils.data.distributed import DistributedSampler
from torch_geometric.loader import DataLoader

train_sampler = DistributedSampler(train_data)
train_loader = DataLoader(train_data, batch_size=16, sampler=train_sampler, num_workers=8, pin_memory=True)

而点属性预测和边链接预测则不同，训练样本一般是从一张大规模图中随机游走采样而成，特别是当需要考虑一些多步特征的情况，分布式训练所需的训练样本拆分就难以操作。

那么这两种任务就无法进行并行训练了呢？幻方AI就这个问题进行了调研和实验，总结了如下两种并行训练方式：

1. 分割子图

一种并行训练的方式是对大规模图进行图分割。

一般来说大规模图数据里，或多或少会存在一些社群分割，例如社交网络、学术网络等。在同一社群中节点彼此比较相似，链接性更高。因此，对于一个大规模图，可以将图根据连通性切割分为若干个子图，然后这些子图上进行分布式数据并行训练。

然而，这种方式需要进行数据的预处理，前期比较耗时。同时，图分割通过某一指标进行切分，实际上人为去除了某些社群间链接的特征，这在某些场景上来说，不可避免会对最终的落地效果造成影响。

2. 并行轮次

另外一种并行训练的方式，是并行轮次。

并行轮次的出发点是：既然不能在每 Epoch 上将训练样本拆分到多张卡上计算，那可以考虑让每张计算显卡单独计算一个 Epoch，合并梯度，减少训练轮次，进而提升训练效率。

以 DeepGCNs 的点预测任务为例：

for epoch in range(args.epochs):
    train_parts = gdata.random_partition_graph(gdata.total_no_of_nodes, cluster_number=args.cluster_number)
    train_data = gdata.generate_sub_graphs(train_parts, cluster_number=args.cluster_number)
    ...

可以看到，在单卡上每 Epoch 都需要在这张大图 gdata 中随机划分出多个部分，然后依此构建多个子图，组成训练样本。这种方式在单卡训练上没有问题，然后扩展到多卡并行会启动多个进程，每轮每个进程都会单独随机采样出训练样本，这使得无法使用 DistributedSampler 在多卡对训练样本进行划分。

因此，索性不对训练样本进行划分，让每进程（或者说每张卡）单独完成一个 Epoch 的训练，然后再同步梯度。

值得注意的是，这种方式虽然简单易执行，但实际效果上是增加了 Batch Size，因此相应的需要增大一些学习率以保证训练效果。另一方面，每个进程都需要从保存在内存里的大图中采样出子图样本，因此对集群的性能要求相应比较高。

实验

这里我们通过实验来验证并行训练的效果。

数据集

幻方AI的数据集仓库集成了 OGB 大规模图公开数据集，通过 3FS 高速存储加速图数据的读取。具体的使用方式如下：

from hfai.datasets import OGB

gdata = OGB(data_name='ogbn-proteins')
gdata.get_data()

for epoch in range(args.epochs):
    # training model ...

如上例代码，我们加载 hfai 数据仓库中的 ogbn-proteins 数据，将这个大图加载入内存。

模型训练

我们测试“并行轮次”的方法，观察模型的收敛速度。首先我们测试单卡上，DeepGCNs 的收敛情况：

之后，我们将实验扩展到8卡并行，测试 DeepGCNs 的收敛：

从上面两次测试结果中可以看到，使用多卡训练下模型收敛的速度更快，因为在同一时间模型喂入了更多的子图数据，梯度同步获得的损失信息更多。这里需要注意，因为子图是通过随机采样出来的，因此在不同的进程中需要使用不同的随机算子，使得喂入的子图是不一样的。从实验可得，并行轮次可以有效提升训练效率，更快模型收敛，进而能直接节省训练的轮次，节省训练时间。

体验总结

综上，图神经网络的并行加速还是切实可行的。采用不同于常规分布式训练的方式，更多的计算资源可以被利用起以加速图神经网络的训练。

综合体验打分如下：

研究指数：★★★★

该模型着眼的是图数据场景，具有很广泛的研究应用场景。
开源指数：★★★★★

数据和代码都已开源，代码逻辑清晰可读性高。
门槛指数：★★★

数据规模中等，模型计算对算力的需求较高，需要多显卡并行加速训练。
通用指数：★★★★

模型是通用的图神经网络模型，在很多图数据上都可以适配应用。
适配指数：★★★★★

依赖简单，很容易与幻方AI的训练优化工具结合，提效明显。

幻方 AI 紧跟 AI 研究的前沿浪潮，致力于用领先算力助力AI落地与价值创造，欢迎各方数据研究者与开发者们一同共建。

萤火跑模型 | 多卡并行助力 BEVFormer 三小时全精度训练

Freja — Tue, 13 Sep 2022 21:00:00 GMT

感知系统是自动驾驶系统中不可或缺的部分，系统中包括许多视觉任务，例如目标检测分割、目标跟踪、语义和实例分割、车道线检测等等。感知系统通过对周围环境状态的准确估计，为预测、规划等下游任务提供重要的依据。其中如何利用车载环视相机采集到的多个视角图像，实现精准的 3D 目标检测和分割是感知系统中最为基础且重要的任务。

2021年9月特斯拉的 AI Day 上，Karpathy 展示了他们的纯视觉自动驾驶感知框架，核心是基于 Transformer 网络结构，在鸟瞰图（BEV，Bird’s eye view）视角下进行感知，紧随其后，学术界也逐渐涌现了基于 Transformer 和 BEV 的优秀工作，其中 BEVFormer 通过构建时空融合网络，在 nuScenes 和 Waymo 的检测比赛上都榜上有名，受到了诸多关注。

然而，最近不少学者和开发者在复现优化 BEVFormer 模型时发现，当使用多显卡加速模型的训练时，随着并行 GPU 数的增加，训练最终的结果反而变差了。这无疑对自动驾驶模型的产业落地造成了不少制约。

幻方 AI 最近对 BEVFormer 进行了整理和优化，分析了其多显卡并行训练精度降低的原因。通过幻方自研的 3FS、FFRecord 对模型训练进行提速；利用 hfai.autotune 进行自动集群参数的搜索，最终在幻方萤火集群上耗时差不多三小时完成了 BEVFormer 的大规模多卡并行训练，并实现了最终精度的复现。

本期文章将为大家详细描述我们的训练优化过程，为研究者和开发者们提供效率提升的思路。

论文标题：BEVFormer: Learning Bird’s-Eye-View Representation from Multi-Camera Images via Spatiotemporal Transformers

原文地址：https://arxiv.org/abs/2203.17270

项目主页：https://github.com/HFAiLab/BEVFormer

BEVFormer 模型介绍

BEVFormer 模型基于可变形注意力机制，融合多视角相机到 BEV（Bird’s eye view）视角进行感知，适用于多种自动驾驶感知任务。该模型采用稀疏注意力机制提高了计算的效率，通过融合时空信息获得了更加准确的感知任务性能。BEVFormer 在 3 月份的 nuScenes 3D 目标检测任务上获得了第一名的成绩，mAP 和 NDS 两个指标均超过以往方法 3 个点以上。在 Waymo 组织的 Camera Track 比赛中，BEVFormer 模型同样获得了冠军。BEVFormer 模型整体架构如下图所示：

BEVFormer 包括三个关键模块，第一个关键模块是 BEV 空间的查询向量（BEV Queries）。

模型设计了 BEV 视角下的网格特征，初始时为每个网格生成查询向量，通过对时序信息和空间信息的查询，将时空特征聚合到 BEV 查询向量的特征中，该特征为 3D 目标检测和地图语义分割的多种自动驾驶感知任务提供了丰富的感知依据。

为什么选择聚合到 BEV 视角？传感器的观测视角为相机视角，如上图左边图所示，自动驾驶汽车往往在车身装备多个摄像机，从而组成环视视角。BEV 视角相当于在车辆正上方 10-20 米的高度俯视进行观察，如上图右边图所示。

相比相机视角，使用 BEV 视角有两个好处，首先下游的规划、决策等任务都需要在车辆所处的 3D 世界下进行，而相机视角是 2D 像素世界，两者存在维度不匹配的问题，因此将多传感器数据融合到 BEV 视角，更加有利于下游模块的执行。其次，不同的物体在 BEV 视角下没有变形问题、视觉遮挡问题、物体重叠问题等，能够使得模型集中于解决感知问题。

第二个关键模块是空间交叉注意力模型（Spatial Cross-Attention）。

BEVFormer 通过 Transformer 结构融合多个视角的图片特征。在每个时间刻，对于每个栅格的 BEV 查询向量，将其所在位置投影到像素平面获得参考点，在参考点周围进行加权特征采样后更新 BEV 查询向量，从而完成空间的特征聚合。

该模块通过 Transformer 实现了从相机视角到 BEV 视角的特征图空间变化。特斯拉公布的方法是将每一个 BEV 查询向量，与各个视角下的所有特征计算注意力，密集的操作导致了过高的内存开销和计算复杂度。而 BEVFormer 通过将 BEV 视角位置投影到像素视角位置，进行对应位置的稀疏注意力计算，融合多视角的同时，大大降低了计算量和内存开销，与此同时，所学习到的注意力初始分布更加集中，梯度的有效性更高，有利于模型后续学习。

第三个关键模块是时序自注意力模型（Temporal Self-Attention）。

每一时刻的 BEV 特征都与上一时刻对齐后的 BEV 特征进行自注意力计算，以保证时序特征的动态更新。在自动驾驶感知任务中，时序信息能够弥补空间中发生的观测限制，并且提供感知所需要的运动信息。具体而言，首先时序信息是空间信息的很好的补充，当物体重叠导致暂时的遮挡、物体过大导致只能获得部分观测时，模型可以融合前后的时序信息获得更好的特征表示，以更加准确的完成感知任务。此外，时序信息可以提供物体的运动状态，例如物体的运动速度、运动方向等等，来帮助预测、规划等下游任务更好的进行。

如上图所示，最左边初始化了 BEV 视角下的栅格图，对于每一个位置，初始化一个 BEV 查询向量。BEV 查询向量一方面通过空间交叉注意力机制将多视角空间特征进行聚合（图中蓝色箭头），另一方面通过时序自注意力机制将历史的时序特征融合到当前时刻（图中绿色箭头），最终获得融合时空特征的 BEV 查询向量，组合得到可以支持下游任务的 BEV 视角特征。

BEVFormer 实验介绍

数据集

BEVFormer 使用 nuScenes 数据集进行试验，nuScenes 数据集是第一个提供自动汽车全套传感器数据的大型数据集，包括了 6 个相机、1 个激光雷达、5 个毫米波雷达、以及 GPS 和 IMU。与另一个常用的自动驾驶领域 KITTI 数据集相比，其包含的对象注释多了 7 倍多，nuScenes 自动驾驶数据集因其数据的质量、规模与难度都比之前数据集有大幅提升，而获得了很多研究者的关注。

检测指标

BEVFormer 在 nuScenes 数据集的上获得了 0.517 的 NDS（nuScenes detection score）、0.416 的 mAP（mean average precision）。

训练时间

BEVFormer Base 模型需要每张 GPU 上 28G 的显存，在单张 A100 上训练完成（24 个 epoch）需要大约 14 天的时间，在 8 张 A100 上训练完成需要接近 2 天的时间。

高昂的训练时间成本，阻碍了科研者和开发者在模型改进的快速迭代。因此幻方 AI 将 BEVFormer 适配到幻方萤火集群，帮助使用者快速有效的进行实验迭代。

BEVFormer 萤火训练

训练时间对比

我们首先保持默认参数，只通过增加并行 GPU 数（相当于增大 batch size），进行模型训练时间和模型性能的对比，结果如下面表格所示。从表格中可以看出，随着并行 GPU 数量的增加，训练 24 个 epoch 的时间逐渐缩短，从最长 14 天 20 小时缩短到 3 小时 12 分中。表格中 ~ 符号代表估算时间，即通过单次迭代的时间对总训练时间进行估算。另外，单卡的训练由于时间过长，只进行时间估算，未统计性能。

模型	并行 GPU 数量	训练耗时	NDS	mAP
BEVFormer Base	1	~ 14 天 20 小时	-	-
BEVFormer Base	8	~ 44 小时 24 分	0.517	0.416
BEVFormer Base	80	6 小时 28 分	0.460	0.365
BEVFormer Base	160	3 小时 30 分	0.385	0.337

增加 GPU 数量，速度变快但性能下降？

从上面表格结果可以看出，在相同的参数训练下，随着并行 GPU 数量的增加（相当于 batch size 增大），训练时间越来越短。但是与此同时，随着并行数量的增加，最终训练出的模型性能越来越差，如下图的变化曲线所示。当我们使用 160 张卡进行训练时，训练模型最终的 NDS 和 mAP 指标与 8 卡时大约相差 15 个点。

为什么增加并行 GPU 数（相当于增大 batch size），会导致训练结果变差呢？

我们探究后发现：这是因为深度学习网络的损失函数是非凸函数，基于凸优化理论的优化算法容易收敛到局部最优解。而当我们增大 batch size，梯度下降的方向更加确定，随机性更小，更容易收敛到附近的局部最优解，而无法达到全局最优解。

如何增加并行的同时保持性能？

随着 batch size 的增大，样本梯度的方差会同比减小，梯度方向更确定，如果要保持增大 batch size 之前的梯度方差，我们可以选择增加学习率。

增加学习率有两方面好处，一方面可以提高收敛速率，因为大 batch size 下，梯度估计更准，增大学习率可以加快模型收敛到最优解。另一方面有助于帮助模型逃出局部最优解，增大学习率可以增大下降步幅，使得模型有机会逃出局部最优解，从而收敛到更好的最优解。

那么如何调整到合适的学习率呢？以往针对大 batch size 的研究提出了下面两种缩放学习率的方法：

sqrt 缩放：Krizhevsky 等人 [1] 提出，当 batch size 增大 k 倍时，学习率应该增大 sqrt(k) 倍。
线性缩放：Goyal 等人[2] 提出了更简单的线性缩放规则，他们认为当 minibatch 大小乘以 k 时，将学习率乘以 k。

我们从实验中发现，按照上述原则调整学习率确实消除了小批量和大批量之间的部分性能差距，但是调整后的模型性能，与小 batch size 的性能还相差甚大。

为什么会有这样的现象呢？文献 [3] 指出，随着 batch size 变得越来越大，得到最佳测试结果所能允许的学习率范围越来越小。也就是说，当 batch size 比较小时，轻微改变学习率对于模型最终的性能影响不大，因此我们可以很容易的找到合适的学习率。但是当 batch size 变大，轻微的学习率变化对最终训练的模型性能影响也随之变大，因此我们需要不断的精细化查找，才能找到合适的参数，这为我们分布式并行训练带来了挑战。

HFai 优化工具实践

为了获得更高的训练精度和更快的训练速度，我们使用幻方 AI 的 Autotune 和 FFRecord 两个工具对训练过程进行优化。

Autotune: 自动搜索集群最佳参数

Autotune 是幻方 AI 针对萤火集群训练平台设计的自动调参工具，可以根据用户设定好的参数范围自动发起集群训练，并统计不同参数下的模型表现。使用时只只需要增加几行代码，就可以使用 Autotune 进行简单的高效的大规模参数调优。

在 BEVFormer 的训练中，Autotune 很好的解决了前面所提到的大规模并行训练时面临的精细化调参的挑战。根据学习率的 sqrt 缩放原则 [1] 和线性缩放原则 [2]，我们可以大致确定在使用上百张GPU进行训练时应该使用的学习率的大致范围，通过 Autotune 我们可以发起多个训练任务来同时对指定范围内的参数进行搜索，快速筛选出最好的参数组合。如下图所示，使用 Autotune 我们可以一键发起数个调参任务：

这样只需要和正常发起一次训练一样的工作量和等待时间，我们就可以得到调参结果。而在使用 Autotune 完成训练后，我们同样可以简单的查看各参数训练结果的排序，便于筛选最优参数：

从上图可见，在使用 160 张 GPU 训练 BEVFormer 时，我们找到的最优参数组合是：lr 为 1e-3、warmup_iters 为 600、min_lr_ratio 为 1e-4，在这组参数下我们取得了 0.490 的 NDS 和 0.389 的 mAP 指标，相比默认参数增加并行的结果，NDS 提高了 0.115，相比 8 张 GPU 时也仅仅下降了 0.027，但训练时间缩短了 14 倍。

FFRecord：加速数据读取

前面的文章我们介绍了幻方高性能数据格式 FFRecord，使用 FFRecord 格式我们可以获得更快的数据读取速度。针对 mmlab 系列开源仓库，我们提供了 mmconverter，用户可以通过简单的三步将 mmlab 代码中使用的数据格式转化为 FFRecord 数据格式，达到适配幻方集群文件系统和提高训练速度的目的。

如下表所示，使用 FFRecord 格式后，不同并行数量任务都获得了一定的性能提升，最终 3 小时就可以完成 BEVFormer 的训练。表格中 ~ 符号代表估算时间，即通过单次迭代的时间对总训练时间进行估算。另外，单卡的训练由于时间过长，只进行时间估算，未统计性能。

模型	并行 GPU 数量	训练耗时	训练耗时（FFRecord）	NDS	mAP
BEVFormer Base	1	~ 14 天 20 小时	~ 13 天 12 小时	-	-
BEVFormer Base	8	~ 44 小时 24 分	~ 39 小时 12 分	0.517	0.416
BEVFormer Base	80	6 小时 28 分	6 小时	0.498	0.391
BEVFormer Base	160	3 小时 12 分	3 小时 10 分	0.490	0.389

下面曲线将使用默认参数的性能曲线与使用 Autotune 搜索到的最佳参数进行了对比。可以看到，使用 Autotune 找到的最佳参数，随着并行数量的增加，性能下降更加平缓，从 1 张 GPU 增加到 160 张 GPU，NDS 和 mAP 都只下降了 0.027，相比默认参数的下降幅度有了很大的提升。

看到这里的你有没有遇到过模型训练久、实验效率低的问题呢？快来申请幻方助力科研的免费 GPU 资源，使用幻方一系列深度学习训练神器，提升你的训练体验吧！

参考文献

[1] Krizhevsky, Alex. “One weird trick for parallelizing convolutional neural networks.” arXiv preprint arXiv:1404.5997 (2014).

[2] Goyal, Priya, et al. “Accurate, large minibatch sgd: Training imagenet in 1 hour.” arXiv preprint arXiv:1706.02677 (2017).

[3] Breuel, Thomas M. “The effects of hyperparameters on SGD training of neural networks.” arXiv preprint arXiv:1508.02788 (2015).

[4] 代季峰对话张祥雨 | 自动驾驶感知新时代！新一代环视感知算法BEVFormer有哪些优势。https://hub.baai.ac.cn/view/19213

hfai.autotune | 一行代码，自动调参

Daylight — Mon, 05 Sep 2022 21:00:00 GMT

调参是深度学习中的必经过程，研究者们往往需要花费大量的时间进行调参，以获取最好效果的模型参数。对于参数规模不大的小模型，手工运行实验进行参数调优还是可以接受的方法，但当遇上大模型时，只凭人工进行参数优化就会变得难以实操，也难以高效利用GPU集群的强大算力。

为了解决这一问题，幻方 AI 基于萤火集群的特性，研发了一套辅助调参工具，Autotune。用户只需要设定好参数搜索范围，在幻方萤火集群上发起训练，Autotune 会自动帮您尝试所有的参数组合，合理分配算力进行测试，最后统计输出不同参数下的模型表现。

设计思路

hfai.autotune 旨在降低调参门槛，改动少部分代码即可实现自动调参。

萤火集群可以通过 hfai 命令或在 JupyterLab 中手动发起的方式来提交训练任务，然而这些方法并不能支持大批量的任务发起。因此在 Autotune 的设计中我们选择 create_experiment_v2 接口，实现一次提交多任务的目标。

用户在使用时设定需要测试的参数组或直接指定参数范围，Autotune 将需要调优的参数进行组合进行网格搜索。这里需要创建多个可以同时训练的任务，在算力允许的范围内运行尽可能多的任务以加速调参过程。

不同任务的训练结果会集中记录在一个日志文件夹中，统计输出时会从这个文件夹中读取所有参数的结果，进行排序后返回。

任务发起

介绍完 hfai.autotune 的设计思路之后，那么其如何使用呢？

只需要加入一行 Python 代码即可：

hfai.autotune.run(config, args, log_dir)

这里必备的三个参数是：

config: 萤火集群训练任务的基本配置文件，参考配置文件格式。
args: 调参参数，在代码中指定。
log_dir: 记录训练结果的文件夹，在代码中指定。

接下来就可以开始调参优化了。具体的，有两种主要的参数搜索模式：

发起固定参数搜索任务
发起参数网格搜索任务

发起固定参数搜索任务

输入参数为列表形式，每个列表内的元素是包含了超参数的一个字典，如下例所示：

import hfai
args = [{"lr":0.1, "bs":256}, {"lr":0.3, "bs":64}]
hfai.autotune.run(config, args, log_dir)

此时，Autotune 会分别使用args列表中的两组参数发起两个调参任务。

发起参数网格搜索任务

输入参数为字典形式，每个超参数通过指定其搜索范围来发起任务。

每个参数的输入可以有两种形式：

对应一个三元组表示参数搜索的起始值，终止值以及搜索步长
对应一个列表表示参数所有可能的取值

如下例所示：

import hfai
args = {"lr":(0.002, 0.006, 0.001), "bs":(64, 320, 64)}
hfai.autotune.run_grid_search(config, args, log_dir)

args = {"lr":[0.002, 0.003, 0.004, 0.005], "bs":[64, 128, 192, 256]}
hfai.autotune.run_grid_search(config, args, log_dir)

此时，Autotune 会分别遍历args中的 lr 和 bs 两个参数的所有取值，并将其组合成最终搜索的参数组。例如在这里 lr 共有4个不同取值，而 bs 也有4个不同取值，那么就会发起总共 4 x 4 = 16 个使用不同 lr 和 bs 组合的调参任务。

本地调试

为了方便用户在本地进行代码调试，Autotune 提供了一键调试模式，set_debug_mode，方便用户在本地运行代码，测试无误后再发至集群进行大规模搜参。使用案例如下：

import hfai
hfai.autotune.set_debug_mode()
args = {"lr":(0.002, 0.006, 0.001), "bs":(64, 320, 64)}
hfai.autotune.run(config, args, log_dir)

在本地调试一组任务时，Autotune 只会使用所有参数组中的任意一组参数发起单个训练任务。默认情况下，Autotune 在非调试模式下运行。

训练代码改动

hfai.autotune 通过尽量少的训练代码改动来支持自动调参功能。在训练代码中只需要在每次训练开始之前调用 Autotune 提供的 get_args 接口来获取任务使用的参数组合，将参数应用到训练配置中就可以进行训练。在训练完成后，同样只需要调用提供的 report 方法就可以将模型测试结果汇报并记录。

如下例所示：

import hfai
args = hfai.autotune.get_args()
lr, batch_size = float(args["lr"]), float(args["bs"])

# ...
# training codes
# ...
# finish training

metrics = {"prec":0.81, "recall":0.79, "loss":0.15}
hfai.autotune.report(metrics)

结果统计

训练结束后，Autotune 会对每个参数训练得到的结果进行汇总统计，按照指定的方式对指标进行排序，返回表现优秀的实验。

import hfai
hfai.autotune.show_result(log_dir, metric="acc")

下图展示了某任务的调参训练结果：

hfai.nn.CPUOffload | 模型训练的显存节省利器

Freja — Wed, 31 Aug 2022 21:00:00 GMT

在深度学习模型的训练中，研究者与开发者们尝尝会碰到显存不足的问题 (OOM, out of memory) ，比如模型参数规模大，或者训练过程中产生的额外开销大，又或者是程序代码的问题，没有足够的显存资源，对我们科研与开发产生了诸多限制。

以往，我们通过代码优化、梯度累计、半精度等等一系列方法，以降低深度学习模型训练的显存需求，然而随着模型参数规模的发展，业务数据越来越复杂，显存需求快速增长，那么除了这些 trick 之外，还有什么既简单、又高效的显存节省方法吗？本期文章介绍的主角， hfai.nn.CPUOffload，给您提供一条不一样的显存节省之路。

那么 hfai.nn.CPUOffload 为什么可以节省显存？如何使用？它和 PyTorch 原始版本有何不同？本期文章将为大家分享 CPUOffload 设计背后的故事，讲述幻方 AI 设计 CPUOffload 的理念，展现 CPUOffload 的性能和便捷。

概述

训练模型的过程中，显存占用主要包括如下三部分

权重矩阵
前向传播的中间向量
反向传播的梯度矩阵

我们以 ResNet50 网络在输入尺寸为（256，3，224，224）时的计算过程作为示例，分析上述三个部分的显存占用。通过 hfai.utils.profile_memory 工具，我们可以很容易的获得模型各部分的细节，如下图所示（由于网络层数较多，只展示部分网络）：

每一行为一个 nn.Module 类型网络的详情信息：

module name：nn.Module 类型网络的名称
type：网络的类型
parameter size: 网络参数量
activation size: 中间变量的 tesnor 大小（不包括参数）
#calls：被调用的次数
input shape：输入的 tensor 形状
output shape：输出的 tensor 形状
peak mem：峰值显存，forward 过程中的峰值显存减去 forward 之前已经占用的显存
forward time：网络 forward 的时间，如果多次调用，则时间累加

最后三行统计了模型的参数总量、中间变量总量、单次迭代前向传播的时间。从上图可以看到 ResNet50 在批大小为 256 时，参数总量为 97 MiB，中间变量为 20974 MiB，中间变量占比较大。而反向传播由于需要保存权重矩阵的梯度和中间变量的梯度，往往需要前向传播的 2 倍或更多的显存空间。

如果我们要减少整体显存占用，不考虑模型的分开存放，有如下两种思路：

第一种思路是使用更小的模型。当数据量小的时候，我们可以使用更小的模型进行拟合；但当数据量大、分布广时，我们需要使用更大的模型，以达到更好的拟合效果。

第二种思路是减小中间变量的显存占用。网络前向传播产生中间变量，反向传播使用中间变量，因此我们可以在中间变量没有被使用的时间，将其移到 CPU 上，以减少显存占用，这个过程我们称之为 CPUOffload。

CPUOffload 原理

中间变量在 PyTorch 中被称为 save tensor, 我们通过 torch.autograd.graph.saved_tensors_hooks 钩子可以在中间变量产生时（前向传播过程中）对其进行打包（pack），在中间变量被使用时（反向传播过程中）对其进行解包（unpack）。一个简单的使用例子如下：

def pack_hook(x):
    return (x.device, x.cpu())

def unpack_hook(packed):
    device, tensor = packed
    return tensor.to(device)

x = torch.randn(5, requires_grad=True)
with torch.autograd.graph.saved_tensors_hooks(pack, unpack):
    y = x * x

y.sum().backward()
torch.allclose(x.grad, (2 * x))

上面的代码是在中间变量产生时，将中间变量移动到 CPU 上，在需要使用中间变量时移动到 GPU 上，从 GPU 移动到 CPU 的过程为 Host2Device（H2D），从 CPU 移动到 GPU 的过程为 Device2Host（D2H）。

PyTorch 提供了 torch.autograd.graph.save_on_cpu 同样实现了此功能。我们可以很方便的对模型中的某些中间变量进行 CPUOffload，一个简单的使用例子如下：

class TorchCPUOffload(nn.Module):
    def __init__(self, module):
        super().__init__()
        self.module = module

    def forward(self, *args, **kwargs):
        with torch.autograd.graph.save_on_cpu(pin_memory=True):
            return self.module(*args, **kwargs)

model = nn.Sequential(
    nn.Linear(10, 100),
    TorchCPUOffload(nn.Linear(100, 100)),
    nn.Linear(100, 10),
)

x = torch.randn(10)
loss = model(x).sum()
loss.backward()

上面的代码定义了一个包括三层全连接层的网络，通过 torch.autograd.graph.save_on_cpu 的方法，第二层全连接层所产生的中间变量，会在不用时移动到 CPU 中，以减少显存占用。

PyTorch CPUOffload 分析

如上述所述，我们可以实现中间变量的 CPUOffload，但在实际优化测试过程中，我们发现如下问题：

第一个问题是 torch.autograd.graph.save_on_cpu 无法自定义 CPUOffload 比例，进入该环境的模型所产生的中间变量，只能全部 CPUOffload。如果想要细粒度的控制 CPUOffload 的比例，则需要开发者自己计算模型各层的中间变量占比，并在代码中增加对应的钩子，以达到按比例 CPUOffload，实现起来较为复杂。

第二个问题是 torch.autograd.graph.save_on_cpu 速度慢，CPUOffload 是以时间换取显存，但是使用 pytorch 默认的版本，显存节省的成本过高。我们以 ResNet50 为例进行速度和显存的测试，具体数据如下表格所示，可以看到在节省百分之十的显存时，模型训练速度慢了将近一倍。

模型	CPUOffload 比例	100 个 iter 耗时	显存峰值
ResNet50	0	32.339 s	22.4 G
ResNet50	0.1	57.461 s	20.3 G
ResNet50	0.5	165.153 s	11.8 G
ResNet50	1	297.812 s	4.4 G

（以上实验在 python3.6、 torch1.10.0+cu113 环境下完成）

`hfai.nn.CPUOffload` 使用方法

为了解决如上的问题，幻方研发设计了 hfai.nn.CPUOffload，与 pytorch torch.autograd.graph.save_on_cpu 的上述问题对应，hfai.nn.CPUOffload 进行了如下优化：

首先，hfai.nn.CPUOffload 提供了 offload_ratio 参数，自动将对应比例的中间变量进行 CPUOffload，使得开发者可以进行更加灵活的选择和更加简单的开发。

另外 hfai.nn.CPUOffload 通过使用独立的 torch.cuda.Stream 进行 H2D 和 D2H，使得整体速度有了较大提升。

以 ResNet50 为例子，这里为大家展示 hfai.nn.CPUOffload 的使用方法和性能。

from hfai.nn import CPUOffload
import torchvision

model = torchvision.models.resnet50().cuda()
x = torch.randn(256, 3, 244, 224).cuda()
with CPUOffload(offload_ratio=0.1, tag='resnet'):
    out = model(x)

hfai.nn.CPUOffload 的使用十分简单，在上面的代码中，通过简单的一行代码，即可实现将特定比例的中间变量进行 CPUOffload，其中 offload_ratio 指定了 CPUOffload 的比例，tag 标签标识了当前 CPUOffload 的模型名称。

另外 hfai.nn.CPUOffload 的使用十分灵活，我们可以通过不同的标签来进行不同模型的 CPUOffload；在不同的迭代中，我们也可以为相同的模型选择不同的 CPUOffload 比例。

`hfai.nn.CPUOffload` 性能对比

我们将 hfai.nn.CPUOffload 与 torch.autograd.graph.save_on_cpu 在相同模型和相同输入的情况下进行了比较，比较结果如下表所示，可以看到在不同的 offload_ratio 下，hfai.nn.CPUOffload 的训练时间都比 torch.autograd.graph.save_on_cpu 更快。在更低的 offload_ratio 下，hfai.nn.CPUOffload 的独立 Stream 能够保证计算过程和 tensor 移动拥有更好的并行性，因此加速更加明显，在 offload_ratio 为 0.1 时达到了 43.0% 的提速。

模型	Offload 比例	显存峰值	100 个 iter 耗时 (pytorch)	100 个 iter 耗时 (hfai)	提速比例
ResNet50	0	22.4 G	32.339 s	32.339 s	-
ResNet50	0.1	20.3 G	57.461 s	32.727 s	43.0%
ResNet50	0.5	11.8 G	165.153 s	136.115 s	16.6%
ResNet50	1	4.4 G	297.812 s	266.686 s	10.4%

看到这里，您是不是也跃跃欲试了呢？我们欢迎广大研究者和开发者安装体验 hfai，来试用这一深度学习训练神器！

hfai datasets | 运转数据入萤火训练的不二法门

High-Flyer — Mon, 29 Aug 2022 21:00:00 GMT

幻方 AI 发布了其沉淀多年的深度学习套件 hfai ，吸引了众多同行研究员和开发者们咨询试用。整个套件的功能较多，而熟悉掌握了这套规则，是能够轻松地调用起平台的算力资源，从而高效完成训练任务的。

为此，我们专门创建了 “hfai 使用心法” 系列专辑，陆续为大家介绍 hfai 一些功能的设计思路和原理，帮助大家更好更快地习得心法，带着 hfai 这套“神功”游刃有余的应对深度学习作业的各项挑战，举重若轻、例不虚发。

之前的文章为大家介绍了如何上传项目代码，构建运行环境和执行任务，一切准备就绪，就差训练数据了。对于大规模的训练数据，上传、管理和部署都是不小的挑战，幻方 AI 研发提供了一套专门为数据集量身定制的“神功”，帮助大家更好的运行大规模的数据，训练出惊艳的模型。

Studio 数据管理页

在萤火集群上的训练数据集，我们都可以通过 Studio 数据管理页进行浏览与使用。如下图所示：

可以看到，幻方 AI 收纳整理了数十种常用的大规模公开数据集，并持续更新中。如果您所需要的数据集在其中，通过 hfai.datasets 方法就可以轻松获取高性能存储支持下的训练数据，进行快速训练。

当然，这些数据不可能满足所有的研究场景。如果您需要上传您自己的数据，可以点击 创建私有数据集。

私有数据集

私有数据集在幻方萤火集群中进行账号级的隔离，您的数据只有您的账号下可以进行访问。

上传的数据需要满足如下条件：

每个文件不小于 256MB （512 x 512KB）
满足条件1的情况下，文件数量尽可能大于100，小于200

这里需要特别注意的是，大量的小文件构成的数据集会对集群存储性能造成严重影响，因此，如果您的数据是这种大量的小文件，我们建议您将数据转化为幻方 AI 自研的 FFRecord 数据格式。

为什么要转换为 FFRecord ?

幻方 AI 自研了一套适合深度学习模型训练场景的文件读写系统 3FS，和一般的文件系统不同，3FS 文件系统有如下的一些特点：

大量打开、关闭小文件的开销比较大
支持高吞吐的随机批量读取

FFRecord 能够充分利用 3FS 文件系统的高效读取性能，其包括如下优势：

合并多个小文件，减少了训练时打开大量小文件的开销，对存储后端更加友好
支持随机批量读取，提升读取速度
包含数据校验，保证读取的数据完整可靠

因此，我们建议用户将训练数据尽可能都转换成 FFRecord 格式，以利用萤火集群的高性能存储，获得极致的训练加速体验。转换指引和案例代码可以阅读这里。

数据上传

用户上传数据的完整流程如下：

在 Studio 中点击 【创建私有数据集】，填写数据信息；
幻方管理员沟通，审核数据内容；
审核通过，点击 【更新授权】 获取上传 token;
使用阿里云OSS工具，通过授权的token，将符合条件的数据上传；
幻方管理员审核，数据流转到萤火集群；
数据在萤火集群准备好后，用户可以在Studio 中获得数据在集群中的地址

更详细的操作请阅读 Studio 中的私有数据集创建指引。

由上述过程描述中可以了解到，用户数据到萤火集群采用如下方式流转：

用户端 —> 阿里云对象存储 —> 萤火集群

幻方 AI 采用特定的阿里云对象存储（OSS）进行用户数据上传的中转，其上传速度取决于用户自己的网络环境。萤火集群和 OSS 之间有专线连接，保证数据的快速流转。

如果您的数据规模巨大，或网络条件有限，可以联系幻方管理员采用其他方式来上传。

hfai python | 任务提交任意所至，萤火训练行云流水

High-Flyer — Fri, 19 Aug 2022 21:00:00 GMT

为此，我们专门创建了 “hfai 使用心法” 系列专辑，分集陆续为大家介绍 hfai 一些功能的设计思路和原理，帮助大家更好更快地习得心法，带着 hfai 这套“神功”游刃有余的应对深度学习作业的各项挑战，举重若轻、例不虚发。

上两个招式为大家介绍了 hfai workspace 和 haienv，其可以帮助用户快速同步本地（个人电脑、个人集群等）工程目录（代码、数据），环境到远程集群中。而这套组合拳下来，其实我们可以认为是“蓄力”和“提气”的过程，接下来就是“神功”最核心的部分，本次文章将为大家介绍 hfai python，其可以帮助大家方便快捷地发起、管理训练任务。

使用场景

如之前文章所介绍的，用户将本地的数据、代码、环境传到远程萤火集群之后，接下来就可以提交任务训练模型了。然而，如何将萤火集群的算力均匀合理的分配给不同任务，用户又如何查看管理自己的任务呢？

这些功能都可以通过 hfai python 来实现。

基本概念

在使用之前，一些基本概念先介绍一下。

幻方AI提出了任务级分时调度的概念来管理萤火集群，算力的分配以任务为单位而不是用户。用户可以自由地提交任务，自由地申请算力规模，由集群调度系统统一分配资源。如下图所示：

计算资源在时间维度上进行切分，将算力分配到不同的任务上。这种方式既提升了集群整体的利用率，也方便了用户对于不同任务弹性伸缩用卡规模，人人都有机会调用大规模的算力进行 AI 科研。

在萤火集群上，训练任务分优先级，高优先级的任务可以抢占低优先级的算力资源。同等优先级内的任务以“先到先服务”的原则提供算力资源。hfai 提供了很多便利的工具（如：hfai.client，hfai.checkpoint），方便您的代码在训练过程中接收集群调度信息，设置断点保存。

本地代码调试

一般的，我们平常执行深度学习训练，例如训练 Bert：

python bert.py -c large.yml

当您需要使用幻方萤火训练模型时，您可以在本地调试修复问题，再提交集群运行。hfai 提供了模拟打断的信号帮助您测试集群运行方式的适配性，如下面例子：

hfai python bert.py -c large.yml ++ --suspend_seconds 100 --life_state 1

这里

++ 标志了当前任务是在本地进行模拟调试
--suspend_seconds 表示每隔多少秒发送一个模拟打断信号
--life_state 表示任务记录的训练位置，这个数值用户可以自己在代码里设置，详情参考这里。

通过如上方式，您可以测试您的训练代码是否可以准确接收集群的打断信号，做好模型的断点保存，同时在下一次任务拉起时，是否可以就当前断点继续往下执行。

提交任务

完成了本地调试之后，您就可以向萤火提交任务了，命令如下：

hfai python bert.py -c large.yml -- --nodes 1 --name train_bert

这里

-- 标志了当前是要提交任务
--nodes 表示使用多少个节点进行训练，这里任务最少分配一个节点8张A100，如果您的任务不需要这么多显卡，可以在代码中进行设定
--name 表示当前任务的命名。

需要注意的是，外部免费用户不用设置优先级 --priority，集群会自动将闲置算力分配出来调度任务。当然，如果您需要不被打断的高优先级训练任务，可以联系幻方洽谈商用。

任务管理

任务提交成功之后，您可以在 studio 页面中查看你任务的运行情况。如下图所示，界面里展示了所提交的任务，任务的 GPU、CPU 使用率，还有集群整体的繁忙程度等。

您也可以获取到任务每时的训练状态，如下所示，方便您更好的管理，调优深度学习模型。

除了在studio页面上可以控制任务的启停、日志查看，您也可以在本地电脑 terminal 当中，通过 hfai 一系列套件工具进行操作。此外，与 hfai python 类似，幻方AI提供了 hfai bash 套件，可以满足用户用bash脚本精细化管理训练任务的需求。有关更多内容可以查阅官方文档。

在减少网络拥塞上，我们的一点实践（一）

Qiao — Fri, 12 Aug 2022 21:00:00 GMT

对于深度学习开发者和研究者们来说，高性能的算力是助力其研究成功的重要武器。对于影响深度学习训练快慢的因素，人们常常容易忽略网络传输在训练提速中的重要作用。特别是在大规模集群，分布式训练的场景中，网络的拥塞可能直接导致GPU算力的失效,就像空有一段段双向8车道的快速路，但如果道路规划凌乱，高速路也只能沦为大型停车场。

本期文章针对网络这个话题，分享一点幻方 AI 就这个方向上的思考和优化。

先聊聊网络拓扑

在幻方萤火集群中，大规模机器学习训练任务对通信和数据读取的性能要求非常高，所以我们使用了目前高性能计算中常用的高速网络来构建节点之间的互联。高速网络的网线和交换机如何排布，这就是我们需要考虑的第一个问题——网络拓扑。

网络拓扑的构建有点像修路，不但要保证任意两个节点之间都能互相访问，还要让从各个节点发出的数据包都能尽可能畅通无阻地到达目的地。

与畅通无阻相反的概念就是拥塞。这里用一个生活中塞车的例子来理解拥塞，假设现在所有车辆都具有相同的最高时速，并且会按照能跑到的最大速度去开。

某个路口有三条车道，源源不断地有车进来但只有一条车道可以出去。如果三条车道的车都遵守交替通行的规则，那么这三条车道的通行效率就只有 1/3。这就是发生了拥塞；
如果三条车道的某一条车道前面有个丁字路口，直行和右转在同一个车道。直行的车辆会通往拥堵的车道，而右转的车辆会前往另一条畅通的车道，但由于右转的车会被直行的车堵住，假设右转和直行的车辆刚好是交替的，那么这条车道的通行效率也只有 2/3。这就是拥塞扩散。

为了尽可能避免塞车的出现，就得有一个好的道路规划。网络拓扑也是同理。

Fattree 是如今大型 AI 集群常用的一种网络拓扑，如英伟达的 Super Pod 等。

Fattree 的原型，如上图所示。多个交换机用一棵树的形式相连，不同的层有不同的交换机数量，不同层之间的线缆的带宽也等比例变化，但每一层的总带宽都相等。这就是一个无阻塞的网络拓扑。我们还是用塞车的问题来理解这个拓扑，每个交换机可以理解成一个路口。我只要把市中心的路口的车道数加得特别多，就不容易堵车。

这个拓扑的路由规则也很简单，因为树的任意两个结点之间都只存在一条最短路径。从 a 点到 b 点的路径直接使用最短路径就可以了。

然而这样的拓扑显然是不实用的，它对网线和交换机的要求太过苛刻。既然没有带宽这么大的网线，我们可以用多根网线组合起来当成一根大的网线用；既然没有这么大的交换机，我们可以把大交换机也拆成多个小的交换机。于是就演化为了下图所示的 Fattree。

相比于 Fattree 的原型，上图所示的 Fattree 引入了一个新的问题：当一个数据包要从其中一个交换机发送到另一个交换机的时候，往往存在很多种可能的路径，如何判断该走那一条呢？如果大量数据包都同时需要通过同一条线路（链路复用），那还是会产生拥塞。根据幻方萤火的经验，链路复用可能会把集群的整体流量降到正常情况的一半，甚至更低。如何通过合理规划流量的路径，减少拥塞的发生和影响范围，这就是路由算法需要解决的问题。

MLNXSM的路由算法

现在的高速网络基本都是用 Mellanox 提供的 MLNXSM 来管理，MLNXSM 提供了针对 Fattree 拓扑结构的路由算法，虽然它的文档中并没有详细介绍，而且代码也没有开源，但众所周知 MLNXSM 是 OpenSM 的魔改版，有兴趣的读者可以通过 OpenSM 的源码来了解其细节。

核心思路

高速网络交换机的路由规则是，每个交换机上都有一个端口转发表，会根据每个数据包的目的地将数据包转发到对应的端口。而路由算法就是要为每个交换机生成这个端口转发表。

OpenSM 的 Fattree 会为每个叶子节点，也就是每张网卡，在 Fattree 的每一层分配一个下行链路，所有到达这一层的去往该目标网卡的流量都经由这个下行链路往下。同时，对这些下行链路做负载均衡，使得每条下行链路上分配到的目标网卡数量尽可能相等。从而减少从上层到下层的流量之间的链路复用。最理想的情况下，如果 Fattree 各层之间的带宽相等，也就是横截面流量之比为 1:1 时，可以完全避免从上到下的链路复用。这样，拥塞就只会在从下到上的流量中发生，拥塞扩散的范围也会变小。

这听起来很不错，但在实际应用 MLNXSM 的 Fattree 算法过程中，我们还是遇到了一些问题。

实际使用遇到的问题

萤火二号实际的网络拓扑并不是一个标准的 Fattree。除去计算机器和存储机器之外，我们还有一部分测试机器、开发机器和用于管理集群的机器等。它们对于网络带宽的要求相对低很多，因此我们使用更低带宽的网卡或是其他方式将它们接入核心网络。但 MLNXSM 的 Fattree 算法并不能很好地处理这一点，它并不会根据机器的应用类型来计算路由。结果就是，由于我们的机器网卡总带宽比 Fattree 的横截面带宽更大，所以 MLNXSM 算出来的路由表常常会把两个以上的网卡分配到同一个链路上，如果这两张网卡恰巧都是负载比较重的网卡，就会产生拥塞。
MLNXSM 的 Fattree 算法产生的路由表会有比较严重的负载不均衡的现象，标号靠前的几个交换机会分到更多的负载。猜测是因为枚举交换机的顺序导致。
Mellanox 交换机有一个很先进的功能叫 sharp，可以把一部分计算在交换机内完成。在实践中我们目前不打开 sharp，但路由算法仍然会给 sharp 节点计算负载。
如果有几个 Fattree 需要互相连接，路由算法的行为将会很难控制。

我们的一些实践

为了解决上述的这些问题，幻方AI自研了一套适配萤火二号的特性的路由算法。

主要思路如下：

将整体的网络在逻辑上切分成标准Fattree的部分和其他部分。
对于网络拓扑中标准 Fattree 的部分，采用上述的 Opensm Fattree 算法类似的思路，给每个叶子交换机上的端口分配一条独立的从上到下的链路，使得从上到下的流量完全没有链路复用。
对于其他部分，使用最短路等图论算法计算转发路径，但同时需要考虑链路的负载均衡。
对于 Fattree 与其他网络之间互联的部分，同样需要在互联线之间做负载均衡，使得每条互联线分配到的目的节点数量相等。
负载均衡的同时，需要根据目标机器类型把单个应用的流量分散到各个链路，最大化单个应用的峰值带宽。
对于不同带宽的网卡，在负载均衡时采用与带宽等比例的权重。
对于 sharp 节点，仅保证其连通性，而不考虑其负载。
类似于互联线和带宽较低的网卡，比较容易成为拥塞的源头。把经过它们的流量集中到部分交换机上，一旦发生拥塞，也能减小拥塞的影响范围。
发生网络故障时，需要把故障交换机上的负载转移到别的交换机上。

在萤火二号上经过长期的生产环境验证，幻方 AI 自研的这套路由算法很好地解决了原来 MLNXSM 路由算法在幻方集群上水土不服的问题，明显提升了集群的网络使用效率。

更优更灵活的路由方案，幻方AI继续持续探索，欢迎各方同行学者一起讨论，加入共建。

JupyterLab TensorBoard Pro：功能完备的 TensorBoard 插件

Aircloud — Mon, 01 Aug 2022 21:00:00 GMT

对于深度学习开发者和研究者们来说，TensorBoard 是帮助理解深度学习运行机理，训练可视化和调试的重要工具。

最近，幻方 AI 开发并集成了 JupyterLab TensorBoard 插件（JupyterLab TensorBoard Pro），帮助大家在使用幻方萤火集群训练模型的同时，也可以方便快捷的使用 Tensorboard 工具进行深度学习训练的调试。

该插件更新修复了社区里原插件的若干问题，并对性能进行了提升。我们对该插件进行了开源，方便用户在没有幻方萤火集群的 JupyterLab 环境的情况下，仍然可以下载和使用。下载安装方式如下

pip install jupyterlab-tensorboard-pro

开源地址：https://github.com/HFAiLab/jupyterlab_tensorboard_pro

研发背景

目前深度学习社区里面已经有了 jupyterlab_tensorboard（前端插件）和 jupyter_tensorboard（对应的后端插件）两个开源插件，然而这两个仓库已经很久没有更新迭代，对于一些外部提交的PR也没有及时合入（实际上，当我们发布 jupyterlab_tensorboard_pro 之后，jupyterlab_tensorboard 合入了一些社区 PR，但 jupyter_tensorboard 仍然存在无法配合较新的 tensorboard 版本使用的问题），基于此我们判断该项目作者已经不在积极地维护对应的仓库。

同时，现有社区的 TensorBoard 插件存在一定的体验问题：

需要同时安装两个 Python 包，并且需要版本对应，增加了使用和维护的成本；
社区版本插件启动 TensorBoard 后端会持续 reload（从磁盘读取文件并更新），这会对 JupyterLab 的体验产生副作用，如果使用了分布式文件系统或者公用集群，甚至还会影响到其他用户。目前社区版本插件无法通过配置来禁用。
功能体验不够完善，例如点击 TensorBoard 按钮之后会阻塞请求导致 JupyterLab 没有响应，并且没有任何提示，同时也对已经创建的 TensorBoard 实例的可视化管理功能也比较缺乏。

基于以上原因，幻方AI基于现有的社区版本二次开发了 TensorBoard 插件，以解决上述问题。在内部使用一段时间后，插件运行稳定并持续获得正向反馈，因此我们将插件开源以帮助更多的开发者和研究者们。

在萤火 JupyterLab 中使用

创建实例

1. 从 launcher 面板创建

您可以登录进入萤火 JupyterLab，从 Launcher 面板点击 TensorBoard 图标。首次点击会进入到一个初始化面板，您可以在该面板内创建 TensorBoard 实例。非首次进入则会直接进入到第一个活跃的 TensorBoard 实例。

2. 通过快捷命令创建

您也可以打开 JupyterLab 快捷指令面板（ctrl + shift + c 唤起），输入 Open TensorBoard。

设置参数

这里，在初始化面板中，提供了两个参数可以进行设置：

Log Dir：日志目录地址，默认是点击 TensorBoard 时当前侧边栏的目录。您也可以手动填写对应目录。这里建议目录尽可能的细化，因为目录内容少会提高初始化的速度。
Reload Interval：TensorBoard 多久对对应目录进行一次重新扫描，默认是 120s。如果您不需要扫描建议及时关闭，日常使用手动 Reload 即可。这里要注意的是：设置 Reload Interval 之后，TensorBoard 后端持续扫描目录会对 JupyterLab 的稳定性和文件系统都产生一定的影响。

选择好参数点击 Create TensorBoard，会同步创建 TensorBoard 实例，这个时候 JupyterLab 后端是阻塞的，请等待实例创建好之后再进行其他操作。

管理实例

创建好实例后，您可以对 TensorBoard 的实例进行管理。目前提供了以下几个功能：

刷新和列表切换：可以切换成其他的实例的 TensorBoard 后端，这个时候不会销毁实例。
独立页面中打开：可以以独立网页 Tab 的形式打开 TensorBoard。
Reload：即重新初始化 TensorBoard 后端，当文件内容有更新时，可以通过此功能载入新的内容（注：TensorBoard 内部的刷新，不会造成 Reload）。
Destroy：销毁实例，会连同前端面板一起关掉。
Duplicate：重新打开一个完全一样的前端面板，此操作会复用 TensorBoard 后端。
New：额外新建一个 TensorBoard 后端，注意事项可以参考上文。

另外，对于我们创建的 TensorBoard 实例，可以在 JupyterLab 的 Kernel 管理面板一同管理，提供跳转至对应实例和删除等功能。

在本地安装使用

JupyterLab-TensorBoard-Pro 支持的最低 Python 版本为 Python 3.6，不过由于幻方 AI 套件建议的 Python 版本为 3.8，因此建议您在 Python 3.8 的环境中使用该插件，从而方便安装其他幻方 AI 套件。

为了兼容多版本的 TensorBoard 和 JupyterLab，插件本身未指定强制依赖版本，需要提前确认本地环境已经安装如下第三方库:

jupyterlab（V3+）
tensorflow（建议 V2+）
tensorboard（建议 V2+）

如果是一个全新的 Python 环境，可以直接 pip install jupyterlab tensorflow tensorboard 安装当前默认版本即可。

通过 pip install jupyterlab-tensorboard-pro 下载安装之后，会默认安装到当前 Python 环境的 JupyterLab 插件列表中，直接启动 JupyterLab 即可使用，使用方式和注意事项与在萤火 JupyterLab 中使用相同。

总结

JupyterLab-TensorBoard-Pro 旨在提供高性能的 TensorBoard 插件服务，帮助您更方便地在萤火集群中调优深度学习训练。我们欢迎大家来试用，并提出反馈。

haienv | 运行环境的影分身术

High-Flyer — Fri, 29 Jul 2022 21:00:00 GMT

最近，幻方 AI 发布了其沉淀多年的深度学习套件 hfai ，吸引了众多同行研究员和开发者们咨询试用。其中，不少用户朋友对 hfai 中一些功能不太熟悉，使用上存在疑惑，或者仅仅使用了其中非常少的一部分功能。

为此，我们专门创建了 “hfai 使用心法”系列专辑，陆续为大家介绍 hfai 一些功能的设计思路和原理，帮助大家更好更快地习得心法，带着 hfai 这套“神功”游刃有余的应对深度学习作业的各项挑战，举重若轻、例不虚发。

上次文章为大家介绍了 hfai workspace，其可以帮助用户快速同步本地工程目录代码到远程萤火集群中。但是在实际的作业中，我们还会碰到另外一个问题——如何构建运行环境？这跟上一招式讲到的“本地-萤火，数据通路”其实是一套组合拳，那么本次“心法”就为大家介绍 haienv，“本地-萤火，环境共享”——快速实现运行环境的集群部署。

使用场景

如上一期文章所介绍的，用户将本地的数据、代码传到远程萤火集群之后，任务代码还需要依赖一些必要的环境才能实际跑起来。然而，集群为了安全起见，采用的是一个封闭的内网系统。

幻方AI经过多年的积累，内部环境已经可以满足很多深度学习训练的需求。使用方式如下，在运行代码的开头加入：

import haienv
haienv.set_env('202111') 

import ...

之后，提交任务到萤火集群上运行即可。

然而就特定场景而言，用户难免会依赖一些特殊的环境。那么，用户如何方便地将自定义的环境从本地（个人电脑，个人集群）推到远程集群中部署呢？

haienv create 应运而生，您可以以此构建您自己的运行环境。

命令结构

haienv create <venv_name> [--py <py>] [--no_extend]

以上就是 haienv 命令的格式，其与 conda 创建环境类似，会生成一个名为 venv_name 的虚拟环境，python 版本号为 <py>（默认为本地电脑里的 python 版本）。

这里要注意的有两点：

--no_extend 表示不扩展本地环境生成新环境，默认是扩展的。如果不指定 --no_extend 那么新创建的环境只可以在本地使用，不能推送到远程。这是因为依赖的基础环境不能被 haienv 所捕获。
在本地使用 hfai init 初始化用户账户信息时，命令里包含了 --venv_path 参数，这个参数在萤火集群中是 /hf_shared/hfai_envs/ 这个路径，在很多用户的本地电脑里基本不会存在。这里如果不指定 --venv_path，可能在venv创建环境时抛出权限问题。这里建议用户在初始化的过程中指定一个存放自己环境的路径，如下所示：

hfai init <token> --venv_path <local_path>

之后，您就可以顺利创建虚拟环境，并通过 haienv list 命令来查看您所创建的虚拟环境了。

环境使用

虚拟环境安装好之后，打开您本地电脑的 ternimal，可以通过如下命令进入该虚拟环境：

source haienv <venv_name>

进入该虚拟环境后，您可以用 pip install 或者 conda install 的方式安装您所需要的依赖。同时，您也可以人工找到 site-package 的目录并手动导入自己想要的包。以上过程和 python 的流行使用方式一致。

当执行代码时，除了在当前激活的环境下运行，您也可以在项目代码中指定这个虚拟环境。使用非常简单，只需要在主函数代码的最开始引入：

import haienv
haienv.set_env(<venv_name>) 

import ...

后续的依赖都会在您所指定的这个虚拟环境中加载。

环境部署

如果您在 hfai Studio 中开发容器里执行的 haienv 创建的虚拟环境，该环境则默认已经部署在集群中。如果您在本地电脑执行的，则需要将环境推送到远程萤火集群进行部署。命令如下：

haienv push <venv_name>

执行该命令时，本地和后台会发生如下交互：

与 hfai workspace 类似，本地与萤火服务层通过 VPN 进行控制交互后，本地环境会先打包发送给OSS，再通过专线导入到萤火集群中。在训练时，模型代码会去搜索加载集群中的这个环境，弹性计算于多个节点中。

这里需要注意的是：

与 workspace 不同在于，haienv 只有 push 动作，没有 pull，环境只能单向上传到集群里；
在本地执行 haienv，必须指定 --no_extend 参数，环境从零构建；在 Studio 开发容器中执行，则可以不加，从集群默认环境中增量安装

为了您的安装方便，我们建议您在 Studio 开发容器中创建您的虚拟环境。

萤火跑模型 | CLIP-GEN 无需文本训练即可文字生成图像

小杞 — Fri, 22 Jul 2022 21:00:00 GMT

随着Transformer模型的发展，近些年多模态模型获得了长足的发展，使得不同任务不同领域可以实现特征的打通，变换出很多新奇好玩的场景。其中非常热门的就是让AI学会看文作图，即文字生成图像，如OpenAI的CLIP模型，其基于带文字的图像数据集上，训练出很惊艳的效果。

然而，收集这些带文字的图像数据集成本非常高。最近字节在Arxiv上发表了一项文本生成图像 (text2img) 的工作，其利用对抗网络GAN改造CLIP模型，使得 CLIP-GEN 可以不依赖带文字描述的图片数据集，直接使用无文本图像数据集进行训练，通过预训练好的 CLIP 模型建立起文本和图像的映射关系。在很多实验数据中表明，它的效果比VQGAN-CLIP要真实，尤其是泛化能力还比不少用大量“文本-图像”数据对训练出来的模型要好很多。

幻方AI最近复现了该项工作，并通过幻方自研的 3FS、hfreduce、算子，对模型训练和推导进行优化。我们在hfai数据仓库中开源了训练数据，模型代码，旨在帮助研究者和开发者们降低研究门槛。

论文标题：CLIP-GEN: Language-Free Training of a Text-to-Image Generator with CLIP

论文地址：https://arxiv.org/abs/2203.00386

模型仓库：https://github.com/HFAiLab/clip-gen

模型介绍

下图是 CLIP-GEN 模型的整体结构：

CLIP-GEN 模型主要由两部分组成，第一部分是一个 VQGAN 模型，用来学习如何把图像编码成一系列的图像标记（image tokens），通过这些图像标记解码还原成一张图片；第二部分是一个 condition transformer 模型，用来学习如何把文字的 CLIP embedding 映射到图像标记（image tokens）中。

训练、推理的过程分为三步：

预训练 VQGAN：输入图像，把图像编码到码本空间，然后再从码本空间解码为图像。经过预训练之后我们就能够把图像表示成一个图像标记（image tokens) 下的离散序列。
训练 condition transformer：通过第一步的预训练，我们已经能够通过图像标记（image tokens) 来生成图片了，在此基础上，我们训练一个 condition transformer，其旨在学习如何把图像的 CLIP embedding 映射到图像标记。训练的过程中 VQGAN 的参数保持不变。
文字生成图片：由于在 CLIP 中，文字和图像共享同一个嵌入空间，我们可以直接把文字的 CLIP embedding 作为 condition transformer 的输入映射到图像标记上，然后通过 VQGAN 的 decoder 来生成图片

数据集

对于 CLIP-GEN 的训练，我们采用了 COCO Caption 数据集，包含 20 万张图文对（训练的过程中没有使用文本）。我们把 COCO Caption 数据集转换成 ffrecord 格式，整合到了 hfai 数据仓库中，可以直接通过以下方式直接使用：

import hfai

dataset = hfai.datasets.CocoCaption(split='train', transform=transform)

有关更多内容，可以访问 hfai 官方文档：https://doc.hfai.high-flyer.cn/index.html

模型训练与优化加速

幻方AI复现了GLIP-GEN模型，并验证其效果。通过幻方自研的 3FS、hfreduce、算子等优化工具，对模型训练和推导进行优化和加速，具体的包括：

hfai ddp：采用 hfreduce 优化多机多卡通信
hfai nn: 重构深度学习算子，提升性能
hfai datasets: 采用高效数据样本格式 ffrecord，充分发挥 3FS 存储带宽性能

下面进行详细描述。

hfai DDP 通信加速

hfai DDP 内部采用了幻方自研的 hfreduce 高性能通讯框架，能有效提升模型的训练速度，使用方法只需要修改一行代码：

# from torch.nn.parallel import DistributedDataParallel
from hfai.nn.parallel import DistributedDataParallel

# ...... initialize model

model = DistributedDataParallel(model, device_ids=[local_rank])

hfai nn 算子加速

为了进一步提升模型训练速度，我们可以使用 hfai.nn 里的高性能算子，相比于 PyTorch 能带来明显的提升，使用方法只需要增加一行代码：

import hfai

model = hfai.nn.to_hfai(model)  # 自动替换为 hfai 高性能算子

使用说明

幻方AI将所复现的模型和优化的方法都进行了开源，统一归集到 hfai 模型仓库 (https://github.com/HFAiLab/hfai-models) 中，欢迎大家来 star。

下载 CLIP 预训练模型：下载 CLIP 后放至 pretrained/clip_vit_b32.pt，该预训练模型来自 OpenAI.
在 COCO 上训练 VQGAN：通过 hfai python 提交任务至萤火集群
```
hfai python train_vqgan.py --ds coco -- -n 1 -p 30
```

在 COCO 上训练 Conditional GPT：通过 hfai python 提交任务至萤火集群

hfai python train_gpt.py --ds coco --vqgan_ckpt /path/to/vqgan/ckpt -- -n 4 -p 30

训练结果

我们来看看训练完成后，一些文本生成图像的效果。

可以看到，不利用带文本的图像数据集，CLIP-GEN 所生成的效果还是非常逼真的。

体验总结

CLIP-GEN 将对抗网络 GAN 用于改造 CLIP 模型，使得 CLIP-GEN 可以不依赖带文字描述的图片数据集，直接使用无文本图像数据集进行训练，这极大降低了数据收集的成本，推动了该领域研究的发展。通过预训练好的 CLIP 模型建立起文本和图像的映射关系，在很多实验数据中表明，CLIP-GEN 的效果比 VQGAN-CLIP 要真实，尤其是泛化能力还比不少用大量“文本-图像”数据对训练出来的模型要好很多。

综合体验打分如下：

研究指数：★★★★

该模型是多模态领域的最新研究成果，降低了数据收集的成本，推动了该领域的发展。
开源指数：★★★

代码没有开源，但所依赖的方法有其他开源版本，容易复现。
门槛指数：★★★

数据规模大，模型适中，适合多级多卡数据并行训练。一般单卡训练难度比较大。
通用指数：★★★★

该方法适用于多模态研究场景，能在很多类似场景下应用。
适配指数：★★★★★

依赖简单，很容易与幻方AI的训练优化工具结合，提效明显。

幻方 AI 紧跟 AI 研究的前沿浪潮，致力于用领先算力助力AI落地与价值创造，欢迎各方数据研究者与开发者们一同共建。

萤火跑模型 | 可变形的 Attention 助力 ViT 优化

Gao Huang — Mon, 11 Jul 2022 21:00:00 GMT

Vision Transformer (ViT) 模型在各类视觉任务上都展现出了强⼤的性能。因其具有较⼤甚⾄增⼤到全局的感受野，ViT 相⽐卷积神经⽹络（CNN）能更好地对⻓距离依赖关系建模，特别是在⼤量训练数据的情况下，ViT 可以轻易扩展参数以达得 SOTA 的实验结果。但是，ViT 里的 Attention 机制也是⼀把双刃剑，⼤量的 key/value 增加了不少计算量，使模型难于收敛，也增加了过拟合的⻛险。

最近来自清华黄高课题组的研究者们对 ViT 模型中的 Attention 机制进行改进，提出了可变形的 attention 机制。研究者们让所有 query 都跟同⼀组 key 和 value 交互，通过对每个输⼊图像学习⼀组偏移量，移动 key 和 value 到重要的位置。这种设计不仅增强了 sparse attention 的表征能⼒，同时具有线性空间复杂度。

该项目利用了幻方AI深度学习训练平台的算力及加速性能，在大量的场景下进行了实验对比，验证了所提方法的优异性能。在前不久闭幕的 CVPR 2022 视觉领域顶级学术会议上，该项工作进入了 Best Paper 奖项的候选角逐。

论文标题：Vision Transformer with Deformable Attention

论文地址：https://arxiv.org/pdf/2201.00520.pdf

模型仓库：https://github.com/LeapLabTHU/DAT

模型介绍

1. 概述

为了避免过量的 attention，现有工作采取了很多的稀疏化 Attention 的办法：

Swin Transformer 设计了滑动窗口机制，每次在窗口内部计算 Attention；
PVT 将 key/value 进行降采样来节约计算的开销

虽然这些方法十分有效，但是它们手工设计的 Attention 模式容易将与任务相关的 key 和 value 信息丢弃。例如 PVT 降采样会损失特征的细节，而 Swin Transformer 则丢弃了单个窗口之外的内容，不同窗口的交互仅依靠滑动窗口间接实现。

理想情况下，Attention 关注的位置应该根据输入变化，key 和 value 应该关注到输入图像中重要的部分，比如目标检测中的目标物体上。在 CNN 的工作中，可变形卷积（DCN）学习偏移的感受野已经表明非常有效。受此启发研究者们尝试探索 Attention 机制中的可变形结构的设计。

然而，直接将 DCN 的机制实现到 ViT 的主干网络中并不是一个简单的问题，因为 Attention 需要计算每个 query 和其对应感受野内所有 key 的相似度，所以这种方案需要缓存每个 query 的 key 和 value。这样一来 Attention 的空间复杂度变为了平方量级，大部分的计算设备都无法接受这样的复杂度。研究者们巧妙地提出让所有 query 都跟同一组 key 和 value 交互，而通过对每个输入图像学习一组偏移量，移动 key 和 value 到重要的位置，最终实现可变形的 attention 机制。如上图 (d) 所示，研究者们提出的可变形 Attention 增强了 Sparse Attention 的表征能力，同时具有线性空间复杂度。

2. 方法

如上图所示，可变形注意力（Deformable Attention）首先将输入的特征图 $x$ 线性映射为 query： $q=xW_q$ 。接下来，初始化一组均匀网格作为采样点的参考点 $p$ ，偏移量由 query 通过子网络 $\theta_\text{offset}$ 产生：

$\Delta{}p=s\tanh(\theta_\text{offset}(q))$ ，

其中 $s$ 为控制偏移量幅度的超参数。接下来，将偏移量加到参考点上得到采样点坐标，使用双线性插值对原特征图进行采样：

$\tilde{x}=\phi_\text{bilinear}(x;p+\Delta{}p)$ ，

其中 $\phi_\text{bilinear}(\cdot;\cdot)$ 表示双线性采样操作，第一个操作数为特征图，第二个操作数为采样点坐标。具体来说，双线性插值的实现如下图所示。

在此基础上，变形后的 key 和 value 分别通过 $\tilde{k}=\tilde{x}W_k$ 和 $\tilde{v}=\tilde{x}W_v$ 得到。常见于 ViT 中的位置编码则通过采样一张相对位置编码表 $\hat{B}_\text{rpb}$ 得到，

$\tilde{B}_\text{rpb}=\phi_\text{bilinear}(\hat{B}_\text{rpb};R(p_q,p_{\tilde{k}}))$ ，

其中 $R(p_q,p_{\tilde{k}})$ 表示计算 query 和变形后的 key 的相对位置。最后多头注意力中的每个头的输出 $z^{(m)}$ 如下式计算：

$z^{(m)}=\text{softmax}\left(\frac{q^{(m)}\tilde{k}^{(m)^\top}}{\sqrt{d}}+\tilde{B}_\text{rpb}\right)\tilde{v}^{(m)}$ ，

其中 $d$ 为每个头的特征维度，与常规 Transformer 中相同。

与 DCN 类似，将设计好的可变形注意力模块加入到 ViT 模型的最后两个阶段，得到 Deformable Attention Transformer，如下图所示。

实验

研究者们在 ImageNet-1K 图像分类、MS-COCO 目标检测和实例分割、ADE20K 语义分割三个重要的视觉任务上进行实验，数据集皆使用幻方 AI 整理提供的数据集仓库 hfai.datasets，详情可以参考：https://doc.hfai.high-flyer.cn/api/datasets.html

运行以上实验需要大量计算资源。研究者们与幻方 AI 合作，采用一系列幻方 AI 自研的深度学习套件，优化模型提升训练速度，在幻方萤火集群上获得了快速的训练迭代。接下来的几个表格展示了最终的实验结果：

ImageNet-1K 图像分类

COCO 目标检测：

ADE20K 语义分割：

可以看到研究者们提出的 DAT 在不同大小的模型上均有很好的表现。

可视化

为了验证 DAT 的有效性，研究者们展示了一些来自 MS-COCO 验证集的样例。

如上图所示，图中橙色圆形表示移动后的部分具有较高 Attention Score 的 key，圆形大小表示累积的 Attention 分数大小，该分数越大，圆形越大。该结果表明 DAT 学习关注到了输入图像中的重要部分。

总结

DAT 方法的提出，进一步推动了 ViT 模型的落地实践。研究者们创造性提出的可变形 Attention 方案，极大降低了 ViT 的训练成本，提升模型的效果。同时，研究者采用幻方 AI 的一系列深度学习自研套件，极大加速了模型的训练过程，获得了大量的实验结果。我们欢迎更多优秀的课题与幻方合作，一道推动 AI 技术的发展与落地。

综合体验打分如下：

研究指数：★★★★

该模型是 ViT 模型的改进研究，推动了视觉 AI 技术的落地实践。
开源指数：★★★★★

数据和代码都进行了完整且规整的开源。
门槛指数：★★★

数据规模大，模型适中，适合多级多卡数据并行训练。一般单卡训练难度比较大。
通用指数：★★★★

该方法能适用于很多视觉场景。
适配指数：★★★★★

依赖简单，很容易与幻方 AI 的训练优化工具结合，提效明显。

幻方 AI 紧跟 AI 研究的前沿浪潮，致力于用领先算力助力 AI 落地与价值创造，欢迎各方数据研究者与开发者们一同共建。

hfai workspace | 打通本地与萤火之间的任督二脉

High-Flyer — Mon, 04 Jul 2022 21:00:00 GMT

最近，幻方 AI 发布了其沉淀多年的深度学习套件 hfai ，吸引了众多同行研究员和开发者们咨询试用。其中，不少用户朋友对 hfai 中一些功能不太熟悉，使用上存在疑惑，或者仅仅使用了其中非常少的一部分功能。如果说我们把 hfai 深度学习套件认为是一部武功秘籍，而您作为使用者若是习得了当中的心法要义，那么您将带着 hfai 这套“神功”游刃有余的应对深度学习作业的各项挑战，举重若轻、例不虚发。

为此，我们专门创建了 “hfai 使用心法” 系列专辑，从本期文章开始，我们会陆续为大家介绍 hfai 一些功能的设计思路和原理，帮助大家更好更快地习得心法。

本次文章为大家介绍 hfai workspace，其旨在同步本地工程目录代码到远程萤火集群中，更丝滑地利用萤火的算力运行您的项目工程。

使用场景

在介绍 hfai worksapce 之前，我们先简单描述其使用的场景。

幻方萤火是一个远程的服务集群。对于 hfai 的用户，他们需要将本地的数据、代码传到这个远程集群中，才可以训练其模型。之前，幻方 AI 提供了 Jupyter 控制台，让用户通过网页手动上传下载，这种方式有诸多不便。

为了解决这个问题，hfai workspace 应运而生，其可以让用户摆脱 Jupyter 操作的不方便，直接连通本地工程目录和远程萤火集群。您可以在您自己环境，如个人电脑、个人集群等，编写代码，调试模型，然后通过一行命令直接将调试好的代码上传远程萤火集群，利用算力来训练您的模型。实现这个过程，只需要您在自己环境里安装 hfai 深度学习套件，通过 VPN 连通集群，即可开始使用。

下面将为大家详细介绍 hfai workspace 的背后原理，使用可以参考官方文档。

整理架构

如上图简图所示，hfai workspace 分为 数据流 和 控制流，控制流主要是发送本地的请求到萤火服务器（Server）中，通过同步服务（Sync）进行本地文件和远程文件的同步控制。数据流使用外部 OSS 储服务，进行数据和代码上传和下载。这里，外部 OSS 存储通过内部代理（Proxy），实现数据的高速传输和安全保障。

为了保障您以及集群整体的隐私和安全，对于不同的用户来说，无论是在 OSS 外部存储，还是幻方集群存储（Weka），数据目录都彼此隔离。当前，幻方萤火集群仅支持从外部指定的 OSS 存储服务中拉取数据，我们配有专线联通以提升传输的速度、稳定和安全。

推送工作区

当您发起 hfai workspace push 命令，推送您的本地目录代码时，在您本地和远程萤火集群之间，会发生如下操作：

请求萤火，验证文件

如上图所示，当发起 push 请求时，本地请求会发送到远程萤火服务（Server），转交给同步服务（Sync），同步服务校验当前的请求，查询集群存储（Weka）中是否存在该工程目录，获取校验和（checksum），将信息返回给本地，与本地目录进行比对。
请求 OSS，准备上传

如上图所示，当完成远程目录和本地目录的校验之后，本地再一次发请求给萤火服务，由集群通过内部代理（Proxy）请求外部 OSS，获取安全的上传链接和校验（STS）。当获取到 OSS 的上传许可后，您的本地工程目录便可以上传了。

这里值得注意的是，待同步的本地工程目录原则上限制为不超过 200G，为更流畅的使用萤火平台，强烈推荐调用幻方的数据集仓库 hfai.datasets。若您所用数据集未被收录，或是私有大规模数据集，您可以联系幻方管理员，hfai 数据集详情请参考 hfai 官方文档
数据上传

如上图所示，本地目录代码和数据开始上传。这里，hfai 会对当前的目录自动进行整体打包，您可以将您的数据和代码整理好，都放在当前这个目录下。注意，数据的上传是增量的，只会对校验后有改动的文件进行打包上传。

目录包上传到 OSS 之后，会通过内部代理（Proxy）下载到萤火集群的存储（Weka）中，这里同步服务（Sync）会负责将包解开，并在集群存储中恢复本地目录设置。

数据传输的进度会通过萤火服务（Server）发送到本地。您可以实时获取到当前推送操作的进度详情。

拉取工作区

除了推送本地目录到远程萤火集群，您也可以通过 hfai workspace pull 命令从远程萤火中拉取您的训练代码，模型保存的 checkpoint，中间数据等信息。该过程与推送基本一致，在进行了文件校验，OSS 请求之后，数据链路如下所示：

数据过程与推送相反，从萤火存储（Weka）的工作区中拉取新增的文件，通过 OSS 传递到本地工程目录中。进度信息会通过萤火服务（Server）实时反馈。

此外，hfai workspace 还提供了 diff, list, download, remove 等命令，它们的执行过程都是上述描述流程的子流程，详细信息可以参考 hfai 官方文档

案例说明

我们以本地一个新的项目工程 WeatherBench 为例子，展示 hfai workspace 的使用。

初始化成 hfai 工作区：

初次推送到远程萤火集群，文件校验后在集群生成工作区 weatherbench ，通过 OSS 将本地工作区打包上传：

您可以在 hfai studio 中看到您的工作区在远程萤火集群中的位置：

当前目录会放置在集群中 hfai/{username}/ 中，按账号进行隔离。当对本地文件 /src/download_era5.py 进行修改，在推送时会展现文件校验结果：

当远程萤火集群完成训练，保存了训练好的模型，可以获取到：

通过 hfai workspace pull 拉取到本地：

萤火跑模型 | AI 驱动的全球精准气象预报

Vachel — Fri, 10 Jun 2022 21:00:00 GMT

现代数值天气预报（numerical weather prediction, NWP）可以追溯到 1920 年，其基于物理原理，整合了几代气象学者的成果经验，是各国气象部门所采用的主流的天气预报方法。这其中，来自欧洲中期天气预报中心（ECMWF）的高分辨率综合预测系统（IFS）模型，在历年的全球天气预报中表现最佳。

最近，在刚刚闭幕不久的 Nvdia GTC 技术年会上，来自英伟达、劳伦斯伯克利国家实验室、密歇根大学安娜堡分校、莱斯大学等机构的研究者研发了一种基于傅里叶神经网络的预测模型 FourCastNet，它能以 0.25° 的分辨率生成全球关键天气指标的预测，这相当于赤道附近大约 30×30 km 的空间分辨率和 720×1440 像素的全球网格大小，与 IFS 系统一致。这项成果使得 AI 气象模型首次能够与传统物理模型 IFS 进行直接比较。

FourCastNet 采用 GPU 训练计算，比传统的 NWP 模型快约 45000 倍，能量节约 12000 倍，使得它能够以很低的成本生成大量的子模型，进行集合预测。实验表明 FourCastNet 极大地改善了概率天气预报的效果，可以在几秒钟内生成对飓风、大气层河流和极端降水等事件的大规模集合预报。

幻方 AI 最近复现了该项工作，并通过幻方自研的 3FS、hfreduce、算子，对模型训练和推导进行优化。我们在 hfai 数据仓库中开源了训练数据，模型代码，旨在帮助研究者和开发者们降低研究门槛。

论文标题：FourCastNet: A Global Data-driven High-resolution Weather Model using Adaptive Fourier Neural Operators

论文地址：https://arxiv.org/abs/2202.11214

模型仓库：https://github.com/HFAiLab/FourCastNet

演示Demo：https://www.high-flyer.cn/hf-earth/

数据集

欧洲中期天气预报中心（ECMWF）提供了一个公开可用的综合数据集 ERA5。作为第五代全球气候再分析结果，ERA5 将其物理模型数据与来自世界各地的观测数据结合起来，形成一个全球完整的、一致的数据集，以小时级到天级不等，提供包括温度、风量、降水、水文、气压等多项全球气象指标数据，供各种气象预报模型学习。官方地址如下：https://www.ecmwf.int/en/forecasts/datasets/reanalysis-datasets/era5

FourCastNet 使用 ERA5 来训练，专注于预报两大气象变量：距离地球表面 10m 处的风速和6 小时总降水量，以对极端天气、自然灾害的预警。该两项指标的预测也被认为是气象领域比较难准确预测的问题之一。

为此，FourCastNet 选择了 20 个相关气象指标，包括几个不同垂直高度的位势高度、温度、风速和相对湿度，一些近地表变量，如地面气压和平均海平面气压以等，作为模型的输入，以挖掘出变量间彼此的关系对风速与降水的影响。具体变量如下：

模型介绍

为了进行 0.25° 分辨率下的全球气象预测，FourCastNet 采用自适应傅里叶神经算子 AFNO，这种神经网络架构是对 Vision Transformer(ViT) 模型地改进，它将混合操作步骤构建成连续的全局卷积，在傅里叶域中通过 FFT 有效实现。通过这样的设计，将空间混合复杂度降低到 O(N log N)，这允许灵活且可扩展地建模跨空间和通道维度的依赖关系。AFNO 的结构如下图所示：

在空间混合步骤中，GFNet 使用跨通道的逐元素矩阵乘法，FNO 采用混合所有通道的全矩阵乘法，而 AFNO 结合 MLP 和软阈值进行逐块矩阵乘法。

对于气象数据，该研究将多种气象变脸按时间整合成 [time, variable, latitude, longitide]格式，进行 AFNO。具体的：FourCastNet 首先将 720 × 1440 lat-lon 网格上的输入变量投影到 2D patch 网格 (h × w)（patch 大小为 p × p，例如 p = 8），每个 patch 表示为一个 d 维 token。然后，将 patch 序列连同位置编码一起馈送到一系列 AFNO 层。如下所示：

FourCastNet 的训练采用递进式，即以 X(t) 作为输入，预测下一步X(t+1)。一次训练输出多步，与真值对比计算 loss。这里论文里采用 RMSE 作为 loss 的评价指标。

对于降水模型，如上图(c)所示，每一步降水的预测实在主干 AFNO 输出的基础上套接一个降水专有 AFNO 进行综合判断，其基于主干网络的预训练结果。这是因为降水这一指标与其他气象指标的概率分布不同，呈现比较鲜明的长尾特征，使得特征空间非常稀疏。对于降水专有的 AFNO，其与主干 AFNO 基本一致，不同在于输出侧，使用 Relu 进行激活，以防止非负输出。

模型复现

幻方 AI 基于 AFNO 和 FourcastNet 两篇论文的描述，复现了 FourCastNet，并采用幻方一系列优化工具进行提速升级，包括ffrecord、hfreduce、hfai.datasets、hfai.checkpoints等功能。

训练样本集

如前面所述，FourCastNet 模型的训练包括pretrian、finetune 和 precipitation 三个部分，对应的数据输入各有不同。因此，我们将原始的 ERA5 数据进行清洗，构建如下样本数据集：

for xt, xt1, xt2, pt1 in hfai.datasets.ERA5.loader():
    out1 = backbone_model(xt)
    out2 = backbone_model(out1)
    out3 = precip_model(out1)

    loss_pretrain = criterion(xt1, out1)
    loss_finetune = criterion(xt2, out2) + loss_pretrain
    loss_precip = criterion(pt1, out3)

    ...

其中xt代表当前t时刻的全球 20 各气象指标数据，xt1代表t+1时刻，xt2代表t+2时刻，pt1代表t+1时刻的降水指标数据。依次进行 pretrian、finetune 和 precipitation 的训练。

目前该样本数据集已整合进幻方数据集仓库中，可以通过hfai.datasets.ERA5调取使用。更多参考hfai 文档。

模型训练

训练分为 pretrian、finetune 和 precipitation 三个部分，每个部分的输入输出各有不同。

pretrain 部分，以 X(t) 作为输入，预测下一步X(t+1)，旨在缩小预测的误差。训练 80 个轮次。

for xt, xt1, _, _ in hfai.datasets.ERA5.loader():
    out = backbone_model(xt)
    loss_pretrain = criterion(xt1, out)

    ...

finetune 部分，在原有主干模型（backbone）的基础上，多预测一步X(t+2)，进一步提升模型预报的精度。训练 50 个轮次。

for xt, xt1, xt2, _ in hfai.datasets.ERA5.loader():
    out1 = backbone_model(xt)
    out2 = backbone_model(out1)
    loss_finetune = criterion(xt1, out1) + criterion(xt2, out2)

    ...

preciptation 部分，原有的主干模型训练完毕，以X(t+1)的信息，预测同时刻的降水P(t+1)。训练 25 个轮次。

for xt, _, _, pt1 in hfai.datasets.ERA5.loader():
    out = backbone_model(xt)
    out_p = precip_model(out)
    loss_precip = criterion(pt1, out_p)

    ...

我们采用hfai.ddp调用 hfreduce 加速训练，使用 64 张 A100 进行数据并行加速，数据和模型采用半精度，最终整体训练耗时在 13~14 小时左右。

更多内容可以阅读我们模型仓库中的开源：https://github.com/HFAiLab/FourCastNet

训练结果

我们以 2018 年 9 月的特大型台风“山竹”为例，展示我们复现的效果：

台风路径

总水汽浓度（Total column water vapour)

可以看到，模型准确预测了未来 3 天的台风行踪路径，风力和降水范围，强度。

体验总结

FourCastNet 将 AFNO 应用于气象领域，首次将 AI 气象的分辨率提升至与物理模型同等水平，使得 AI 气象模型应用于真实场景中可能性大幅提高。FourCastNet 在大风、极端暴雨的预测表现，为进一步提升天气预报的准确度，预防极端自然灾害提供了新的方案。幻方 AI 对 FourCastNet 的复现与优化，使得该项研究成果切实落地，为环境治理、灾害预警贡献一份力量。

综合体验打分如下：

研究指数：★★★★★

该模型首次将 AI 气象分辨率和准确度提升至物理模型水平，促进了 AI 气象模型的落地。同时，与气象领域知识的融合让该方向有更多研究空间。
开源指数：★

代码没有开源，ERA5 数据下载比较繁琐。
门槛指数：★★★★

数据规模大，模型适中，适合多级多卡数据并行训练。一般单卡训练难度比较大。
通用指数：★★★

该方法是基于 AFNO 的应用，在高分辨率数据场景中皆可以使用。
适配指数：★★★★★

依赖简单，很容易与幻方 AI 的训练优化工具结合，提效明显。

幻方 AI 紧跟 AI 研究的前沿浪潮，致力于用领先算力助力 AI 落地与价值创造，欢迎各方数据研究者与开发者们一同共建。

hfai.checkpoint | 并行训练中的快速保存与加载

小杞 — Mon, 23 May 2022 21:00:00 GMT

分布式训练中模型的保存，特别是大模型，常常需要耗费很多的时间，降低了整体的 GPU 利用率。针对这类问题，幻方 AI 进行了攻关，优化过往深度学习模型单机训练保存的方法，研发出分布式 checkpoint 方案，大幅度降低模型保存与加载上的开销。

分布式 checkpoint

当我们进行分布式模型训练，特别是在训练大模型时，保存 checkpoint 需要较长的时间，这不仅浪费集群的计算资源，并且给集群整体的任务调度带来管理成本。为此幻方 AI 提供了一个分布式保存 checkpoint 的功能。

该功能的基本原理是：假设有 N 块 GPU，我们把模型参数和优化器参数切分成 N 个部分，然后每个分布 rank 把对应的部分写入文件系统；在读取的时候我们从文件系统中读出所有 checkpoint 拼接出完整的模型参数和优化器参数。除了模型参数和优化器参数，其他的信息会由 rank 0 进行保存。

使用方法：

from hfai.checkpoint import save, load

model = MyModel()
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
others = {'epoch': epoch, 'step': step+1}
save('latest.pt', model, optimizer, others=others)

state = load('latest.pt', map_location='cpu')
epoch, step = state['epoch'], state['step']
model.load_state_dict(ckpt['model'])
optimizer.load_state_dict(ckpt['optimizer'])

集群上实际测试性能如下：

可以看到，随着并行程度的增加，使用hfai.checkpoint进行模型保存加载的耗时越来越少。

自动断点训练

对于幻方萤火集群上的训练，需要接收集群统一的调度信号进行训练任务的管理。这里幻方 AI 提供了一个 hfai.checkpoint.init 函数帮助用户进行断点训练，该函数会自动加载上次保存的模型、优化器等状态，返回上次训练的 epoch 和 step，我们可以通过 epoch 和 step 进行优雅断点训练。同时，我们会向 model 注册一个成员函数 try_save ，通过这个函数可以在打断训练之前保存训练的状态。

使用方法：

model = MyModel()
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)

start_epoch, start_step, others = hfai.checkpoint.init(model, optimizer, ckpt_path='latest.pt')
for epoch in range(start_epoch, epochs):
    for step, (x, y) in enumerate(dataloader):
        if step < start_step:
            continue

        output = model(x)
        loss_fn(y, output).backward()
        model.try_save(epoch, step, others=None)

通过上述封装，您可以在代码中省去很多断点操作，简单方便地将代码适配幻方萤火系统，进一步降低门槛。

萤火跑模型 | 高精地图构建模型 HDMapNet 助力更精准的自动驾驶

Freja — Mon, 09 May 2022 21:00:00 GMT

高精地图是自动驾驶系统的关键模块，可以有效提升自动驾驶汽车的行驶安全度，强化自动驾驶系统的整体感知能力和决策能力。然而传统的高精地图构建流程复杂，需要消耗大量的资源和人力，导致其扩展性和实时性无法满足自动驾驶的使用需求。

CVPR 2021 Workshop 最佳论文提名的一篇文章 HDMapNet，来自清华大学交叉信息院 MARS Lab，提出了一个使用深度神经网络实时构建局部高精地图的思路，该方法是从传统的高精地图构建流程转向端到端方法的关键工作。

幻方AI最近对这项工作进行了整理和优化，在幻方萤火二号上复现了实验。通过幻方自研的 3FS、hfreduce、算子，对模型训练进行提速，整合代码、简化接口，融入hfai数据仓库和模型仓库中，帮助研究者和开发者们降低使用门槛。本期文章将为大家详细描述。

论文标题：HDMapNet: An Online HD Map Construction and Evaluation Framework

原文地址：https://arxiv.org/abs/2107.06307

项目主页：https://tsinghua-mars-lab.github.io/HDMapNet

模型仓库：https://github.com/HFAiLab/hdmapnet

模型介绍

高精地图

高精地图分为影像地图和语义地图，自动驾驶使用的通常是高精地图的语义地图部分，因此本文的高精地图特指高精语义地图。高精地图是车道级别、厘米精度的道路网络和交通信息地图，包括且不限于如下标注信息：

厘米级车道边线和中心线几何信息
车道线类型（白虚线/白实线/黄虚线/黄实线）
车道线类（高速/城市道路/自行车道）
交通信息（红绿灯/限速）

高精地图作为自动驾驶汽车的高精度“记忆”模块，是自动驾驶决策的关键依据，离开了“记忆”，无论眼睛（摄像头和雷达）和思考（控制决策系统）有多么发达，都无法对事件有全局的把控。高精地图为自动驾驶系统提供了：

厘米精度的定位
超过相机视野的地图记忆
反应当前路况的实时信息

传统高精地图构建流程

构建传统的高精地图需要经过一系列流程：

使用车辆到道路上采集信息，需要大量的人力和车辆资源。
将收集到的数据进行处理，包括点云拼接、地图生产等。
对点云中的信息进行识别，例如道路、车速牌等。
人工验证识别结果的准确性，确定最终的标注。

构建完成后还需要定时进行更新维护，以适应不断变化的路况信息，这样繁杂的流程需要消耗大量的资源和人力，导致其扩展性和实时性无法满足自动驾驶系统的使用需求。

HDMapNet

HDMapNet 提出了一种使用车载传感器（相机、雷达）来动态构建局部高精地图的方法，主要创新点如下：

将高精地图使用向量化的方法进行表示，包括：
- 语义分割向量化：使用三个类别表示道路边缘、道路分割线、斑马线三类线段
- 实例分割向量化：使用不同的编号标识不同的线段
- 方向识别向量化：使用 36 个集合表示 360 度的道路方向
使用深度神经网络学习从相机前视图到鸟瞰视图的转换

HDMapNet的整体网络结构如上，输入为 RGB 图片或者雷达点云，两种输入可以单独使用，也可以同时使用，实验结果表示，RGB 图片和雷达点云同时使用时能达到最佳性能。

RGB 图片通过基于 EfficientNet 的图像编码器进行编码，如下面图中所示，本车周围六个方向均匀放置了六个相机，在每一个时刻，将六个相机获得的 RGB 图片编码成六个前视图的隐含表征。视角变换步骤（Neural view transformation）利用相机外参，将前视图的隐含表征转换到鸟瞰图的隐含表征，如下面图片的中间步骤到右边步骤的对应关系。最后使用鸟瞰图解码器从鸟瞰图的隐含表征中解码得到最终的鸟瞰图表示。

雷达点云先进行像素化，再通过点云网络得到鸟瞰图的隐含表征，最后通过鸟瞰图解码器解码得到最终的鸟瞰图表示。

将 RGB 输入和雷达点云输入处理得到鸟瞰图表示后，HDMapNet 使用三个分支：

语义分割预测：预测鸟瞰视角下，道路边缘、道路分割线、斑马线三个类别，使用二元交叉熵作为损失函数。
实例分割预测：用不同编号标识不同线段，使用 DiscriminativeLoss 作为损失函数。
道路方向预测：用 36 类分别代表 360 度的道路方向，使用二元交叉熵作为损失函数。

模型实践

幻方 AI 基于 lift-splat-shoot 复现了 HDMapNet，再基于 HDMapNet 开源代码进行修正，采用幻方一系列优化工具进行提速升级，包括 hfreduce 并行训练、hfai 数据集接入、hfai 集群训练断点挂起等功能。

数据集

HDMapNet 使用 NuScenes 数据集进行试验，NuScenes 数据集是第一个提供自动汽车全套传感器数据的大型数据集，包括了6个相机、1个激光雷达、5个毫米波雷达、以及GPS 和 IMU。与另一个常用的自动驾驶领域 KITTI 数据集相比，其包含的对象注释多了7倍多。

这里幻方 AI 将 NuScenes 数据集进行处理，将原始数据集转换成 ffrecord 的格式，用户调用简单的接口便可以加载高性能的 dataloader 进行模型训练。具体如下示例：

from hfai.datasets import NuScenes

dataset = NuScenes(split='train')
loader = dataset.loader(batch_size=64, num_workers=4)

for data_dict in loader:
    # training model ...

仅仅需要上述 4 行代码，我们便可以获得如下格式的 NuScenes 数据集：

模型训练

登录幻方萤火二号，引入 hfai
初始化 hfreduce 分布式参数
对于每一轮训练，加入接收集群调度的逻辑代码，并做好模型 checkpoint 的保存

训练结果

我们申请 4 个节点，32 张 A100 显卡进行并行训练。通过hfai命令行工具提交任务，训练效果如下图所示：

在萤火二号上训练 HDMapNet，每个 Epoch 耗时仅 5 分钟左右，相比单卡性能提升了 32 倍以上。GPU 利用率在 95% 左右，充分利用其了显卡的计算性能。模型在第 24 轮基本达到了收敛状态。

训练效果可视化如下：

可视化效果1：
可视化效果2：

体验总结

HDMapNet 新颖的提出了使用深度神经网络构建局部高精地图的方法，减少了传统构建方法耗费大量人力和资源的问题、增加了高精地图的可扩展性和实时性，有助于构建更加智能的自动驾驶系统。幻方 AI 对 HDMapNet 的实验进行了复现，借助萤火平台对 HDMapNet 的训练进行了优化，最终达到了超过单卡 32 倍的训练速度。

综合体验打分如下：

研究指数：★★★★

该模型以低成本、端到端的方法实时构建局部高精地图，是自动驾驶领域的重要课题。
开源指数：★★★★★

数据处理和代码都已经开源，代码逻辑清晰、可读性高。
门槛指数：★★★

数据量、模型大小中等，普通高性能单卡即可运行。
通用指数：★★

该方法是基于神经网络构建高精地图方法的早期工作，地图丰富度、定位精度、预测性能有待提升。
适配指数：★★★★★

依赖简单，很容易与幻方AI的训练优化工具结合，提效明显。

幻方 AI 紧跟 AI 研究的前沿浪潮，致力于用领先算力助力AI落地与价值创造，欢迎各方数据研究者与开发者们一同共建。

开始使用 hfai

High-Flyer — Tue, 26 Apr 2022 21:00:00 GMT

幻方萤火 AI 训练平台 hfai 可以为您提供简单易用的深度学习训练服务。使用 hfai，您可以在幻方萤火二号上自由部署训练大规模深度学习网络，使用幻方萤火的优化工具为您的模型训练加速。

观看以下视频快速了解如何通过控制台使用 hfai。

官方文档：https://doc.hfai.high-flyer.cn/index.html

宣传推文：https://mp.weixin.qq.com/s/qslsTUeq06bu09IVJ5TwBA

案例代码：https://github.com/HFAiLab/hfai-models/tree/main/resnet

使用说明

您可以申请使用，审批通过后就可以获得 hfai 专属账号和资源。您通过 VPN 接入幻方萤火的网络，浏览器访问 https://studio.yinghuo.high-flyer.cn，通过预设账号登录即可使用。

使用样本数据格式 ffrecord

ffrecord 是 hfai 提供的适用于模型训练场景的样本数据格式，其与幻方萤火集群的结构相适应，能提供高性能的批次数据读取。更多信息，请参见《幻方萤火 | 高性能数据格式 ffrecord》。

使用并行训练工具 hfreduce

hfreduce 是 hfai 提供的高速模型并行训练工具，根据幻方萤火集群的特性专为计算节点设计的 All Reduce 工具。其使用 CPU 做加法运算以计算总梯度。更多信息，请参见《幻方萤火 | 模型并行训练工具 hfreduce》。

使用深度学习算子 hfai.nn

hfai.nn 是 hfai 中的算子模块，其提供了比 Pytorch 原生算子性能更优异的优化算子，助力一些常见场景下的模型训练加速。更多信息，请参见《幻方萤火 | 性能卓越的深度学习算子 hfai.nn》

萤火跑模型 | 气象预测之时空图神经网络

Vachel — Wed, 09 Mar 2022 21:00:00 GMT

气象预测直接关乎人类社会的生产生活，一直以来是人类重点研究的科学领域之一。传统的天气预报采用非常复杂的数值气象模型，近年来越来越多的机器学习，深度学习方法被应用到这一领域，他们将气象预测抽象成一个时空预测问题，尝试通过多维时间序列、图神经网络（GNN）等方法来解决。 AAAI 2022前不久公布了今年的 Accept Paper List，其中来自西湖大学的几位研究者联合发表了一项科研成果。他们针对近年来交通流预测任务下的图神经网络模型，认为其架构不能直接照搬进气象预测领域并进行了证明。提出了一系列改进方法，最终模型取得了SOTA的效果。

幻方AI团队最近复现优化了该项研究成果，来看看最近火热的GNN是否可以有效提升气象预测的精度，缩短预测时间。

论文标题：Conditional Local Convolution for Spatio-temporal Meteorological Forecasting

论文地址：https://arxiv.org/abs/2101.01000

论文源码：https://github.com/BIRD-TAO/CLCRN

模型介绍

在交通流预测领域中，“GNN+RNN”的方法近年来获得了许多成功，其静态切片空间的图结构抽象，到动态时间上的embedding累积，有效解决了时空预测场景下的图动态问题。与交通流预测类似，气象预测也可以抽象为一个时空预测问题。然而，只是简单的套用交通流预测的方法，往往不能取得很好的效果。

气象预测领域的特殊性有：

数据的不规则性：分布世界各地、太空卫星的各种气象传感器所捕捉的气象数据往往是不规则的，使得经典的CNN无法适用。
非平面的数据信号：气象数据通常是立体的球形信号，而非平面信号；陆地和海洋的温度传感器各有不同，不在固定的网格结构中，但气象预测往往无法将其割裂开来单独做分析。
高度的时空依赖：不同地形表现出完全不同的风流或温度转移模型，而极端气候事件往往使得数据不符合平稳性条件

如上图所示，气候预测任务与交通流预测任务的不同在于，每一个局部点受其四面八方的因素影响，局部空间下模式应该基本类似，因为气象学中，热和风是自由扩散的。而交通流预测任务中，两个靠近的交通枢纽之间模式因为交通流向会存在很大差异。

因此，研究者们需要找到一种方法，来抽象气象预测场景下的问题。具体的：

将分布地球表面的各种气象信号抽象为一个球数据结构，给予坐标，球体上要区分若干个局部空间（local space），他们与实际地理测地线、欧氏距离相关；
通过单位速度测地线（unit-speed geodesics）定义出每个局部空间的邻居，进而构建起等轴映射（isometric map），将球体数据结构转化成图数据结构，以描述不同局部空间之间的关系；

针对上述问题抽象，该文研究者们提出了**因地制宜的图卷积核（location-characterized kernel）**来改良DCRNN（一种用于交通流预测的图神经网络，ICLR 2018）中的图卷积部分。具体的：

如上图a，为了能更好的保留地球表面地理标线的相对方位，作者采用圆柱切线空间（cylindrical-tangent space）和水平映射（horizon map）方法构建起局部空间并映射其邻居，构建起等轴映射；
为了更好了模拟气象学中的局部一致性，作者重新设计了条件局部卷积核，以满足因地制宜的卷积特征，邻近局部特征相似和地理特性不同下的相邻卷积核共享三个条件。如上图b，来自东方的xp对西边的两个局部空间 i 和 j 的影响应该是一样的，不会因为xp离 j 近就在图卷积中被赋予更大的信息量；
如上图c，两个相邻空间各自拥有不同的地理分布，这时不能简单的采用平滑方法以满足邻近特征相似。作者提出结合其距离和角度来重新计算卷积核的权重。具体的，在计算测地线距离的基础上，使用极坐标重新刻画数据，计算相邻点的角平分线。

最后，模型整体的结构如下图：

模型实践

作者开源了实验代码，包括baseline，其主要采用图神经网络库torch-geometric进行进一步的封装实现。代码由于是刚刚开源，需要一定改造方可运行起来。目前数据和代码已整合进幻方AI的数据仓库与模型仓库当中，供后续研究者们调用。

1. 数据集

该研究采用的是全球天气预报基准数据集 WeatherBench，由大数据地球研究社区 Pangeo 于2020年开源。其一共由2048个分布在地球各个重要位置上的统计节点。作者选取温度、湿度、云量、地表风分量四个天气预报项进行试验，以小时为单位划分数据，以过去12小时数据进行分析，对未来12小时进行预报。

值得注意的是，在开源的代码中，作者是将这四个预报项单独进行预测，分别测试准确度，而不是综合多个因素进行分析。例如，对温度的预测只考虑过去温度的变化，不会将过去湿度，云量，风量等因素综合考虑其中。

除了时序化的指标数据，作者还抽析出了每个指标每个时刻的地理位置信息，以构建一个球体结构。如上一章所介绍的，这个结构将被解析成一个动态图，进行时空分析并预测。所抽析的地理指标包括：经度、纬度、海拔、陆地标记。最终，组成的输入输出形式为：

2. 模型训练

作者开源的代码采用的单机单卡训练模式，提供了预处理好的训练样本与测试样本，可以通过pickle直接加载使用。在幻方萤火二号上，用户可以获得更多的高性能显卡，因此可以进行并行训练以获取更好的性能，相关内容可以参考之前的文章《PyTorch分布式训练方法》。

这里，我们采用分布式数据并行（DDP）的方式改造代码，提高训练性能。具体的

用 DistributedDataParallel封装原始模型

from hfai.nn.parallel import DistributedDataParallel

self.model = DistributedDataParallel(clcrn_model.cuda(), device_ids=[local_rank])

用DistributedSampler封装原始数据集，使其自由分散在多张显卡中

from torch.utils.data.distributed import DistributedSampler

train_datasampler = DistributedSampler(trainset)
self.trainloader = torch.utils.data.DataLoader(trainset, batch_size, sampler=train_datasampler, num_workers=4, pin_memory=True)

加入hfai监听集群中的统一调度信号，做好checkpoint保存

import hfai

# the first node gets the suspend signal, save checkpoint
if rank == 0 and hfai.receive_suspend_command():
    self.save_model("checkpoint.m")
    time.sleep(5)
    hfai.go_suspend()

通过Process托管多进程，管理多张显卡上的训练

import torch

if __name__ == '__main__':
    torch.multiprocessing.spawn(main, args=(), nprocs=torch.cuda.device_count())

改造完成后，运行的效果如下：

主节点Epoch平均耗时48s左右，相比单机单卡训练提升速度6-7倍。完整的改造代码已集成进幻方AI的模型仓库中，感兴趣的研究者们可以滑至文末申请体验，了解更多的技术细节。

类似WeatherBench中展示的可视化案例，可以将预测出的天气信息映射到地图中，对比出天气的变化情况,。如下图所展示的温度变化案例：

体验总结

气象预测与人类生产生活息息相关。AAAI 2022这篇研究工作将近年来火热的时空图神经网络应用到了气象预测当中，将这一问题具体抽象成了图数据结构并予以解决。该工作进一步推进了AI+环境科学的发展。同时，无论是模型还是该场景中的问题，还有许多能进一步研究的空间。比如，是否可以结合更多源的数据对某一预报项进行预测、不同因素之间彼此是否有关联、扩大数据规模增加算力能提升多少精度等等。我们欢迎感兴趣的学者加入我们，继续在幻方AI团队的萤火平台上对该课题进行进一步探索。

综合体验打分如下：

研究指数：★★★★

气象预测是重要的研究课题，结合AI方法有很多的创新空间。
开源指数：★★★★

数据和代码都进行了完整开源，代码部分修改后可以运行。
门槛指数：★★★

作者开源的数据规模不大，对算力的需求较小。但随着数据规模的提升和多指标的整合分析，算力的需求会逐渐增大。
通用指数：★★

模型针对气象预测领域的特性进行模型设计，通用性比较差。
适配指数：★★

模型采用单机单卡方式实现，需要进行多机并行训练改良方可适配幻方萤火二号。

幻方 AI 紧跟 AI 研究的前沿浪潮，致力于用领先算力助力AI落地与价值创造，欢迎各方数据研究者与开发者们一同共建。

Alphafold 训练优化 03 | 踩坑日记

Daylight — Tue, 01 Feb 2022 15:00:00 GMT

前两期文章展示了幻方AI对Alphafold的优化，采用特征预处理和特征裁切两种方式提高了Alphafold数据处理性能，通过并行训练加速神器进一步提高模型的训练速度，将Alphafold深度融合进幻方AI的集群特点中，发挥最大的计算效能。

那么从整体上看，在幻方萤火二号上训练Alphafold还有哪些需要我们注意的，以及未来同类型深度学习模型该如何优化？就这些话题，本期文章将和大家聊聊幻方AI的一些思考。

模型仓库：https://github.com/HFAiLab/alphafold-optimized

总体优化效果

首先，我们看一下在进行了前两期优化后，最终Alphafold在幻方萤火二号上的训练效果。我们使用16个节点共128张A100显卡，对将近7万条蛋白质序列进行了比较完整的Alphafold训练，得到了如下的实验结果：

平均GPU利用率	90.14%
平均每Epoch时间（min）	84.15
平均每次迭代时间（s）	11.32
每Epoch开始时数据加载时间（s）	197
总迭代数	35680
总训练时长（h）	112

训练时的GPU利用率曲线如下：

训练时的GPU利用率大多数时候在90%左右。可见经过上述的多种优化后，Alphafold 已经能够在幻方AI平台上高效进行训练，而有中间偶尔出现的利用率骤降则是训练开始时的数据加载耗时。

另外，从CPU利用率图可以看出，整个训练过程中CPU的利用率长期保持在大于50%（35/64核）的情况，说明Alphafold训练时确实对CPU算力有非常大的需求。

多卡训练中的一些踩坑

在优化Alphafold模型的过程中，特别是将单机单卡拓展至多机多卡时，我们遇到了不少的问题。

1. OpenFold代码问题

Deepmind没有开源Alphafold训练代码，现在github主流的Alphafold训练代码是由其他三方实现的，比如OpenFold（https://github.com/aqlaboratory/openfold），其通过PyTorch比较完整地复现了Deepmind开源的Alphafold 2（https://github.com/deepmind/alphafold）。

OpenFold的复现使用了比较复杂的代码构造模型结构，其中不同蛋白质序列数据在训练时可能会有不完全一样的前向传播路径，这导致在多卡训练时，一些机器上有一部分参数没有参与梯度计算，而Pytorch DDP提供的find_unused_parameters方法并不能正确的找到这部分参数，因此会时不时报错，无法正常进行多卡训练。

经我们一段时间问题排查发现，是Loss Function部分中存在部分不参与计算的Loss，需要将其手动注释掉后方能正常进行多卡训练。

2. 环境问题

另外一个比较容易踩坑的就是环境问题。幻方AI平台内置了PyTorch等环境，适配了常见的深度学习机器学习。而Alphafold本身作为一个深度学习领域的新方向，所依赖的环境必然与常规环境不同。

因为幻方AI训练平台采用的是计算存储分离的方式，环境统一配置，不托管镜像，所以Alphafold的运行环境需要一步步重头安装，这里不可避免地会发生与原有环境版本冲突，需要补充系统库等一系列问题。

因此，我们通过 haienv 完全创建一个全新的虚拟环境，并安装了 Alphafold 的依赖项。然而，在进行多机多卡训练时会出现NCCL Error的问题。经过排查后发现：由于 haienv 创建了全新的环境，使用了与原先不同的全新Python bin文件，而新建的Python并没有获得IPC Lock，因此在多机多卡训练时NCCL会出现问题。需要在sudo权限下使用setcap指令setcap cap_ipc_lock+iep /venv/openfold38_0/bin/python来解决问题。

最后，我们将Alphafold所依赖的环境单独列出来，供用户调用。

局限与展望

虽然经过我们的优化后模型本身训练的效率已经很高，GPU利用率达到90%以上，但我们仍然对模型性能做出了一定的取舍，并且模型内部仍然存在不少优化的空间。一方面来说我们已经完成的优化都基于不修改模型内部结构的前提，因此对模型内不同部分的差异造成的额外开销难以处理；另一方面来说训练时显存的利用还不够理想：A100 40GB的显存只能使用到21GB而有接近一半的空余，但batch_size只能设置为1，这造成了大量浪费。

从以上两个角度入手，后续Alphafold模型还可以进行的优化有以下三个方向：

结构简化：Alphafold 2本身虽然达到了很好的蛋白质预测准确率，但其复杂的模型结构和极高的资源占用使得复现和部署模型推理都有很高的门槛。因此，仍然可以考虑对模型内部进行细致的Profile，寻求对模型结构进行一定程度的简化。
算子优化：幻方AI在算子优化，尤其是Transformer类模型的算子优化方面有很多的积累，而Alphafold作为一个同样基于Transformer结构的模型，还没有用到这些算子优化技术。因此，在这里也有进一步优化的空间。
模型并行：Alphafold目前训练时的显存占用率仍然不高，单卡batch_size也只能设置为1。可以考虑尝试在Alphafold上应用更多的模型并行方式，例如和Pipeline结合来增大训练时的单卡batch_size，提升显存使用效率。

欢迎各方数据科学家与开发者们一同共研共建。

Alphafold 训练优化 02 | 多卡训练提速

Daylight — Fri, 07 Jan 2022 15:00:00 GMT

上一期文章讲到，幻方AI通过优化数据处理，采用特征预处理和特征裁切两种方式提高了Alphafold整体的训练性能。众所周知，幻方AI有很多并行训练加速神器，比如hfreduce，3FS，hfai.nn算子库等，它们是否能对Alphafold整体的训练进一步加速呢？本期文章将就这些问题进行试验。

hfreduce

之前的文章《幻方力量 | 模型并行训练工具：hfreduce》提到过，由于幻方AI架构的特点，Nvidia官方提供的NCCL工具并不能充分发挥萤火二号集群的通信带宽，因此幻方AI自研了hfreduce工具优化显卡间的通信，提高多机多卡的并行训练效率。使用方式非常简单，只需要将pytorch DDP替换为hfai DDP即可：

model = hfai.nn.parallel.DistributedDataParallel(model, device_ids=[local_rank])

对于Alphafold模型，其本身的特点是前向传播路径比较长，而反向传播路径较短，具有比较鲜明的“参数规模不大，但计算量大的特点”。我们采用数据并行将数据分散到不同的显卡上进行并行加速，以加快训练效率。因此我们为Alphafold的训练引入了hfreduce，利用hfai.ddp工具进行数据并行，尝试进一步推高Alphafold在萤火二号集群上的计算效率。

然而在实践中我们发现，hfreduce给Alphafold训练带来的训练提升在绝对值上并不大。这是什么原因呢？在随后的实验部分中我们结合实验结果来分析这一问题。

模型仓库：https://github.com/HFAiLab/alphafold-optimized

实验设定

我们设计并进行了三组实验，以对比使用hfreduce后的加速效果。实验的设定如下表所示

序号	GPU数	DDP Backend	迭代次数	耗时统计对象	特征裁切
1	1	无	200	loss.backward()	最大
2	128	NCCL	200	loss.backward()	最大
3	128	hfreduce	200	loss.backward() + model.sync_grad()	最大

我们希望通过多卡和单卡训练的耗时对比，来直接看到多卡通信带来的额外开销。因此，我们对3种训练方式测试了其200次迭代的平均backward耗时。

其中，“特征裁切：最大” 代表了对特征进行了尽可能大的裁切，目的是避免因特征处理导致多卡训练时部分GPU上的梯度同步长期处于等待状态，影响结果的准确性。使用hfreduce时由于梯度同步被单独拆出，因此需要将两个函数合并统计梯度计算+梯度同步的总时间。

此外，训练的基本参数还包括：

训练样本数量：73472
Batch Size: 1

实验结果

根据上文中的实验设定，我们测试NCCL并行加速和hfreduce并行加速的完整训练结果如下表所示：

	NCCL	hfreduce
每个step平均时长（s）	11.27	11.26

直接对比单次迭代的总耗时，从结果来看在Alphafold的训练场景下HFReduce和NCCL的不同选择并不会带来显著的效率差异。理论上来说，HFReduce比NCCL更加适合萤火二号的集群架构，一定能带来比DDP更快速的梯度同步。但在Alphafold的实验结果中，两次实验中每个Step的迭代速度几乎完全相同，这是为什么呢？

对于这种现象，最主要的原因可能是Alphafold独特的模型结构和蛋白质数据的特点使得模型虽然需要很大量的GPU运算，却只需要传输很少的梯度数据。Alphafold的模型参数大小只有93.2M，模型的参数量不到1亿，使用FP32单精度的情况下占用显存372.8MB。在前向传播时Alphafold会将模型中的主要部分（Evoformer）复制4次，这会带来一个很长的前向传播通路，使得时间复杂度大大增加。然而，最终模型需要更新的参数依然只有原先的数量，因此显得梯度同步的开销在整个迭代中显得微不足道了。

那么，既然梯度同步的耗时不显著，我们是否可以通过一些profile方法拆分出每个部分的耗时，进行定量优化呢？我们对backward部分进行更细致的测试，最终测试结果如下表所示：

序号	GPU数	DDP Backend	Backward平均耗时（s）	通信开销（s）
1	1	无	5.63	-
2	128	NCCL	6.03	0.4
3	128	hfreduce	5.91	0.28

这里，不考虑多卡时因数据处理进度不同产生的等待，那么多卡训练时的backward时间减去单卡训练的backward时间既可以认为是多机多卡带来的额外通信开销。

上表的结果中可以看到，在萤火二号上使用DDP训练时，选择hfreduce能够比NCCL节省30%的通信开销，这是一个非常可观的数字，说明hfreduce确实更适合在幻方萤火二号上进行多机多卡训练的加速。同时，90%以上的耗时在forward的计算上，特别是layer_norm和attention的算子计算开销比较大，这里需要进一步使用hfai的优化算子进行进一步提升。

hfai.nn

如前文《幻方萤火 | 性能卓越的深度学习算子 hfai.nn》所提到的，幻方AI针对Pytorch框架进行了深度优化，结合萤火二号的集群特点，对一些常用的AI算子重新设计研发，进一步提升了模型整体的训练效率。其使用非常简单，只需要在原有的任意模型训练代码中加入一行：

model = hfai.nn.to_hfai(model)

hfai会自动扫描您的代码，替换为优化后的算子。hfai.nn提供了当前主流深度学习模型中常用的MultiHeadAttention，LayerNorm，LSTM等结构的优化算子，能够大幅加速模型中这些算子的运算。我们尝试在Alphafold中引入hfai.nn算子库，测试其能够对模型训练产生多大的增益效果。

Alphafold GPU开销分析

由于hfai.nn的算子加速取决于模型中各类型运算占总GPU开销的比重情况，因此我们首先尝试使用Pyorch Profiler工具对Alphafold模型使用标准torch算子时的GPU运算时间进行一些分析。在一次总耗时为12.6s的迭代中，主要开销情况如下：

Name	CUDA Time	CUDA %
aten::native_layer_norm	2.105s	16.81%
void at::native::(anonymous namespace)::RowwiseMomen…	1.866s	14.90%
aten::mm	1.716s	13.70%
aten::native_layer_norm_backward	1.642s	13.11%
atten::add_	1.454s	11.61%

仅前5类运算占了Alphafold训练开销的70%，而其中LayerNorm与Attention中的矩阵运算又是最主要的耗时来源：仅LayerNorm运算就占了总耗时的30%左右。Alphafold相对一般的BERT等Transformer类模型更为复杂，使用了自己实现的Attention，因此在注意力运算上无法获得hfia.nn的加速。但同样占据耗时大头的LayerNorm却能够使用到hfai算子加速，因此可以预期hfai.nn能给alphafold带来较好的加速效果。

为了探究具体加速效果，我们首先尝试对Alphafold中LayerNorm的使用情况进行一些分析。一般来说输入Tensor的形状往往会对算子加速的效果产生较大的影响，因此在这里我们首先尝试对输入模型中的LayerNorm层的不同的Tensor形状进行一些理论上的分析。在一个典型Transformer类模型中，LayerNorm的输入形状往往是[BatchSize, SeqLen, EmbDim]。在Alphafold中由于其模型结构与简单的Transformer模型有较大差异，输入的张量形状会更为特殊一些。模型中LayerNorm层的不同Input Shape对应的出现频率可见下表：

Input Shape	Run%	Relative Perf%	Torch Run Time(s)	Hfai Run Time(s)
[1,256,256,128]	59.96%	364.07%	0.003	0.0008
[1,132,256,256]	13.2%	267.02%	0.0017	0.0007
[1,1,256,256,64]	8.03%	379.80%	0.0028	0.0007
[1,128,256,256]	7.92%	276.52%	0.0017	0.0006
[1,256,132,256]	4.4%	276.13%	0.0017	0.0006
[1,256,128,256]	2.64%	272.92%	0.0017	0.0006
[1,256,384]	1.87%	80.50%	0.0002	0.0003
[1,1024,256,64]	1.32%	471.77%	0.0111	0.0024
[1,256,1024,64]	0.44%	474.66%	0.0111	0.0024
[1,4,256,256,64]	0.11%	474.41%	0.0111	0.0024
[1,256,256]	0.11%	93.81%	0.0003	0.0003

可见在大多数输入时，hfai.nn提供的LayerNorm算子都能取得相比torch.nn数倍的大幅性能提升，只在极少数形状下性能可能会稍差于torch。另外根据上表中不同输入Tensor形状的占比可以计算得，使用hfai.nn在Alphafold中可以在LayerNorm的GPU开销上取得总计350%的加速。由于Alphafold中的LayerNorm使用频率很高，且在模型的不同部分都有使用，因此可以预计对模型整体也能有较明显的加速效果。

实际训练加速结果

在理论分析了hfai.nn能给Alphafold带来的加速幅度后，我们也希望能了解在实际训练时Alphafold能够从hfai.nn中获得的收益。因此我们在与Alphafold预训练时相同设定的真实场景下使用128卡进行了并行训练，分别测试了使用torch算子和hfai算子进行Alphafold训练时的训练耗时情况，结果如下表：

算子	GPU数	迭代次数	平均时长
torch.nn	128	200	11.26s
hfai.nn	128	200	8.64s

相比于使用torch.nn算子进行训练，使用hfai.nn时模型整体训练中单次迭代的平均时长能够从11.26秒减少到8.64秒。可见在Alphafold训练时只需要添加一行代码引入hfai.nn，就足足能够获得30%的训练性能提升。值得注意的是，由于Alphafold中自定义了Attention实现没有使用标准的nn.MultiHeadAttention，在这里能够获得的hfai.nn算子加速效果其实只来自于LayerNorm运算减少的耗时。由此可见在常见的标准Transformer类模型中，使用hfai.nn将更加容易获得大幅的性能提升。

总结

通过上述两个实验，我们验证了 hfreduce 和 hfai.nn算子对alphafold模型的加速效果。他们分别表现出了不同的加速能力，说明在模型优化的过程中我们需要根据模型的结构特性选择合适的加速方案。这里，幻方AI将持续加大研发，提供更多的优化方案和工具。

萤火跑模型 | CLIP 模型

小杞 — Tue, 28 Dec 2021 21:00:00 GMT

CLIP（Contrastive Language-Image Pre-Training，以下简称 CLIP）模型是 OpenAI 在 2021 年初发布的用于匹配图像和文本的预训练神经网络模型，可以说是近年来在多模态研究领域的经典之作。该模型直接使用大量的互联网数据进行预训练，在很多任务表现上达到了目前最佳表现（SOTA）。

本次我们尝试使用 Google 开源的 Conceptual Captions 数据集来训练 CLIP 模型，并对其效果进行一定的验证。

本次体验接入的系统是幻方AI团队研发的幻方萤火二号：自研智能分时调度、高效存储与网络通信，可以实现任务级别的弹性算力支持；配合API或者 JupyterLab 交互就可以轻松接入，利用起澎湃算力进行大规模深度学习训练。

源码地址：https://github.com/openai/CLIP | https://github.com/mlfoundations/open_clip

数据集：https://ai.google.com/research/ConceptualCaptions/download

模型仓库：https://github.com/HFAiLab/clip-gen

CLIP 模型简介

在目前图像分类领域中，我们训练的模型通常会遇到以下问题：

模型需要用到大量的格式化标注数据，这些标注数据获取通常成本高昂。
模型在当前数据集的效果比较好，但是可能模型的泛化能力较差，同时迁移到新的训练任务也比较困难

与此同时，互联网上面已经存在了大量的图像文本对（在网页中，开发者一般都会为图片添加一段文字备注），实际上这些素材可以作为已经标注好的数据集，利用这些数据集进行训练，既能解决获取标注数据成本高昂的问题，同时也因为互联网上的数据量比较大和数据本身差异较大，更容易让我们获得泛化能力较强的模型。

CLIP 模型就是基于上述概念，使用 OpenAI 收集到的 4 亿对图像文本对，分别将文本和图像进行编码，之后使用 metric learning 进行训练，通过计算 cosine similarities 计算图像和文本的匹配，通过最大化匹配的图像文本对的 cosine similarity 和最小化不匹配的文本对的 cosine similarity 来优化目标函数，它的核心流程比较容易理解，可以直接参考下述伪代码：

# image_encoder - ResNet or Vision Transformer
# text_encoder  - CBOW or Text Transformer
# I[n, h, w, c] - minibatch of aligned images
# T[n, l]       - minibatch of aligned texts
# W_i[d_i, d_e] - learned proj of image to embed
# W_t[d_t, d_e] - learned proj of text to embed
# t             - learned temperature parameter
# extract feature representations of each modality
I_f = image_encoder(I) #[n, d_i]
T_f = text_encoder(T)  #[n, d_t]
# joint multimodal embedding [n, d_e]
I_e = l2_normalize(np.dot(I_f, W_i), axis=1)
T_e = l2_normalize(np.dot(T_f, W_t), axis=1)
# scaled pairwise cosine similarities [n, n]
logits = np.dot(I_e, T_e.T) * np.exp(t)
# symmetric loss function
labels = np.arange(n)
loss_i = cross_entropy_loss(logits, labels, axis=0)
loss_t = cross_entropy_loss(logits, labels, axis=1)
loss   = (loss_i + loss_t)/2

在预测阶段，也是通过一系列生成的文本对和目标图像，计算 cosine similarities 从而获取预测值：

CLIP 通过以上过程，取得的效果还是比较惊艳的。

OpenAI 官方开源了 CLIP 模型部分的代码，不过如果想完整的的复现训练，还需要写不少训练相关的代码，这里我们基于 github 上的另外一个开源实现，通过一定的修改，来完成在萤火平台上训练的整个过程。

数据预处理

下载和清洗

Conceptual Captions 数据集是 Google 从数十亿互联网网页提取的图像文本信息数据，并进行了若干种类型的过滤，使得数据集具有较高的质量和准确度，总共约 300 万个测试数据和 8000 多个验证数据。

它提供的是描述和图片的下载地址对，其数据格式如下：

title	filepath
a man on a scooter , taking part in the rally	http://l7.alamy.com/zooms/5e84d0b73a544fe8a8e5ac0d20052685/a-man-on-a-lambretta-scooter-taking-part-in-the-daily-express-rally-ddmmff.jpg
the dentist drill the tooth with a turbine	http://l7.alamy.com/zooms/da764b930a6e477da792ddb8b83b1ecb/the-dentist-drill-the-tooth-with-a-turbine-hn3aph.jpg
...

我们可以将该数据集通过 python 脚本下载和统一存放，下载完之后，我们需要将该数据集进行清洗，将下载的空图片或者不完整的图片数据过滤，之后才可以使用。

这部分下载和清洗，我们简单编写 python 脚本即可完成，我们完全下载清洗后得到的 csv 文件格式为：

title	filepath
a man on a scooter , taking part in the rally	/demo/path/clip/cc_data_1/val/29/29.jpg
the dentist drill the tooth with a turbine	/demo/path/clip/cc_data_1/val/19/19.jpg
...

实际上，这个数据集由于比较大，而且都是分散在各个 URL 而不是一个完整的压缩包，我们下载会花费不少时间，累计接近 100G 的解压大小也会对我们的数据存放造成一定负担，不过幻方AI团队目前正在基于自研的高性能平行文件系统进行数据集仓库的建设，也就是说，对于一些常用的数据集和已经在集群上训练过的公开数据集，包括以上 Conceptual Captions 数据集，我们会经过处理后放在可访问的公开文件目录，后续在萤火使用相同的数据集进行训练或模型开发，就无需重复下载和预处理数据集了，直接使用即可，从而能极大减轻我们的数据负担。

在分布式训练开始之前，虽然我们已经可以使用清洗好的图像数据进行训练了，但如果我们频繁打开小文件，对训练自身的性能和整个集群的影响都会比较大，我们可以使用幻方AI自研的 FFRecord 来将多个小文件进行合并处理，从而减少训练时打开大量小文件的开销，同时对存储后端更加友好，我们会在下文对其进行介绍。

FFRecord 数据格式转换

FFRecord（FireFlyer Record）是幻方AI自研的简单高效的存储二进制记录的文件格式，它的一些特点包括：

合并多个文件，减少了训练时打开大量小文件的开销，对存储后端更加友好。
支持随机读取，可以适应不同的样本读取模式。
包含数据校验，保证读取的数据完整可靠。

同时，为了方便训练，我们还提供了专为 PyTorch 和 FFDataset（FireFlyer Dataset）和 FFDataLoader（FireFlyer DataLoader），优势包括：

高效读取，使用 Linux Asynchronous I/O 的接口，充分利用了自研存储系统的随机读取性能优势。
简单易用，只需要对使用原有 PyTorch DataLoader 的代码进行简单的修改即可切换为 FFDataLoader。
使用灵活，完全兼容 PyTorch Dataset 和 DataLoader 的相关接口，如 Sampler 等。

数据预处理和 FFRecord 转换部分的代码，均可在萤火二号访问。

在转换为 FFRecord 之后，我们就可以开始调整开源代码，进行训练了。

开源代码调整

这里我们针对开源代码进行调整，这里的主要调整的内容包括：

参数调整，包括超参数、文件路径等
入口文件微调，引入 haienv、hfai
使用 FFDataset 代替原有的 Dataset
适配集群打断逻辑

step1: 参数调整部分比较简单，我们可以根据自己的需要调整数据集和验证集的文件路径，以及 batch-size 和 epoches 等参数。

这里我们增加了两个参数，分别用于表示训练集和验证集的 FFRecord 文件地址：

'--train-ddr-perfix', '/3fs-jd/prod/private/nxt/clip/cc_data_1_proc/train',
'--val-ddr-perfix', '/3fs-jd/prod/private/nxt/clip/cc_data_1_proc/val',

step2: 引入 haienv、hfai 的代码和我们的上一篇 informer 模型的调整比较类似，这里我们直接给出代码

import haienv
haienv.set_env("202111")

import os
import sys
sys.path.append(f'{os.getcwd()}/src')

step3: 为了提高训练性能，我们使用上文提到的 FFDataset 代替 torch 的 Dataset，这里我们给出一个实现代码：

class FFDataset(fftorch.Dataset):
    def __init__(self, input_filename, transforms, img_key, caption_key, sep="\t", ddr_perfix=None):
        logging.info(f'[FFDataset] Loading csv data from {input_filename}. use ffrecord: {ddr_perfix}')
        df = pd.read_csv(input_filename, sep=sep)

        self.images = df[img_key].tolist()
        self.captions = df[caption_key].tolist()
        self.transforms = transforms
        self.ddr_perfix = ddr_perfix
        super().__init__(self.ddr_perfix, check_data=True)
        logging.info('[FFDataset] Done loading data.')

    def __len__(self):
        return len(self.captions)

    def process(self, indexes, data):
        '''
        在每次读取样本数据的时候，
        FFDataset 先从 FFRecord 文件中读取样本的二进制数据，
        然后将原始二进制数据传递给 process 函数。
        用户需要继承 ffrecord.torch.Dataset 并自定义 process 函数以对原始二进制数据进行处理。
        '''
        samples = []

        try:
            for bytes_ in data:
                sample = deserialize(bytes_)
                image_from_bytes = sample['image_bytes']
                images = self.transforms(image_from_bytes)
                texts = tokenize([str(sample['caption'])])[0]
        
                samples.append((images, texts))
        except:
            raise

        return samples

step4: 由于萤火的任务调度规则为分时分优先级调度，也就是说任务提交以后，有可能会被集群打断暂停，任务打断实际上会强行结束当前进程，后续再启动的时候就又会再次执行初始化逻辑，为了能够在打断后恢复继续训练，我们需要及时保存我们的 checkpoint，并且在任务恢复的情况下加载之前的历史记录继续训练，同时，我们应该在整个训练结束之后对模型进行保存。

这里笔者提供一个比较简单的方式，我们利用 whole_life_state 来进行标记，第一次 whole_life_state 是 0，如果是打断恢复的场景，这个时候 whole_life_state 就是 1，因此我们加载最新的 checkpoint 进行继续训练，这里给出关键的代码：

# main 函数里面处理 whole_life_state:
whole_life_state = hfai.get_whole_life_state()
print('whole_life_state:', whole_life_state)

if whole_life_state == 1:
    # 只要不是第一次运行，我们就尝试加载最近的 checkpoint：
    args.resume = os.path.join(args.checkpoint_path, f"epoch_latest.pt")

hfai.set_whole_life_state(1)

# 加载 checkpoint:
if args.resume is not None:
    if os.path.isfile(args.resume):
        checkpoint = torch.load(args.resume, map_location=loc)
        start_epoch = checkpoint["epoch"]
        sd = checkpoint["state_dict"]
        model.load_state_dict(sd)
        if optimizer is not None:
            optimizer.load_state_dict(checkpoint["optimizer"])
        logging.info(
            f"=> loaded checkpoint resume:'{args.resume}' (epoch {checkpoint['epoch']})"
        )
    # ...

经过以上处理，我们就可以验证和运行我们的任务了。

测试运行并提交训练

幻方萤火提供的 hfai 命令行工具提供了测试运行的接口，我们可以通过 hfai python src/training/main.py 来进行测试运行（一般来说，我们建议测试运行没有明显问题之后再提交训练）

测试运行会在当前终端直接打印运行信息：

2021-12-22,22:34:14 | INFO | Rank 0 | Use GPU: 0 for training
2021-12-22,22:34:14 | INFO | Rank 4 | Use GPU: 4 for training
2021-12-22,22:34:14 | INFO | Rank 6 | Use GPU: 6 for training
2021-12-22,22:34:14 | INFO | Rank 3 | Use GPU: 3 for training
2021-12-22,22:34:14 | INFO | Rank 7 | Use GPU: 7 for training
2021-12-22,22:34:14 | INFO | Rank 1 | Use GPU: 1 for training
2021-12-22,22:34:14 | INFO | Rank 2 | Use GPU: 2 for training
2021-12-22,22:34:14 | INFO | Rank 5 | Use GPU: 5 for training
... ...
Train Epoch: 0 [0/2850879 (0%)]	Loss: 8.067907	Data (t) 42.464	Batch (t) 72.875	LR: 0.000000	logit_scale 2.659
Train Epoch: 0 [3840/2850879 (1%)]	Loss: 8.055143	Data (t) 0.008	Batch (t) 0.356	LR: 0.000001	logit_scale 2.659

我们可以看到，我们的代码已经可以在单机8卡的场景下测试运行了，不过，我们会发现由于数据集比较多，这个训练还是比较慢的，此时我们可以利用萤火的多机多卡进行并行训练。

我们在幻方萤火的管理页面提交任务运行，选择 6 个节点：

任务提交后，我们在幻方萤火的管理页面，可以比较方便地看到任务的日志和状态：

我们可以发现，在 6 节点（48 张显卡）并行训练的过程中，单个 epoch 训练大约 6 分钟，我们如果训练 30 个 epoch，加上初始化时间，大约 3 个多小时:

验证

针对 CLIP 模型的验证测试，openAI 提供了多种方式，我们可以比较方便地在萤火上对此进行复现，这里为了便于我们理解模型，我们展示一种比较可视化的方式对它进行简单的验证:

我们使用 skimage 的一些图片和文本信息，使用我们训练好的模型计算相似度，然后通过二维矩阵的方式打印出来：

我们把相似度最高的，标记成了黄色，我们可以从上图看出，我们训练的模型在该情况下表现尚可。

总结

CLIP 模型思路新颖同时比较简单，借助萤火集群和自研文件系统，我们可以比较高效地训练该模型，在训练的过程中，GPU 的平均使用率达到 80%+, GPU Memory 使用率持续在 10GB+。

基于目前的数据和模型，我们后续还可以进行进一步的探索。

综合体验打分如下：

研究新颖度: ★★

该模型对图像文本分类领域提出了一个较为新颖的方法，有一定的技术推进性。
开源指数: ★★★★

官方开源了模型代码，并且 github 上有一些更加通用的开源代码参考，代码逻辑清晰可读性高。
算力需求: ★★★

解压后训练数据大小为 84 G，近 300 万条训练数据，6 节点 48 张 A100 大约需要 3 小时
通用指数: ★★★

该模型虽然比较容易理解，但是在多模态领域还是比较有启发的，有可能会在一定程度上成为后续深入研究的 baseline
模型适配度: ★★★★

额外依赖不多，在使用下载好的数据集的情况下我们只需要简单改一下 DataSet 等代码即可运行

幻方 AI 紧跟 AI 研究的前沿浪潮，致力于用领先算力助力AI落地与价值创造，欢迎各方数据研究者与开发者们一同共建。

Alphafold 训练优化 01 | 数据处理优化

Daylight — Sat, 25 Dec 2021 15:00:00 GMT

若说2021年人工智能学术界最令人振奋的成果，那么Alphafold可谓当之无愧。Alphafold2在CASP14 蛋白质预测挑战上取得了远远超出同类模型的准确率，并首次将蛋白质结构预测的精度提高到了原子级别——已经接近了实验测量的水准。

幻方AI团队在Alphafold2推出不久就在萤火二号平台成功将Alphafold2训练运行了起来，详情如我们上一期文章《萤火跑模型 | Alphafold 蛋白质结构预测》。然而，由于Alphafold2本身模型的复杂性和开源代码的缺陷，并且幻方AI的训练平台也比较特殊，使得一开始复现资源利用效率非常底下，训练成本不可接受。

为了提高Alphafold的训练效率，帮助科研人员与开发者们降低研究门槛，幻方AI团队对Alphafold2模型进行了大量优化，大幅提高了模型训练的效率。我们的工作可以被概括为以下三点：

重构代码，解决开源版本中无法多机多卡训练的问题，利用数据并行大幅提高Alphafold2训练效率，降低复现门槛；
优化数据处理逻辑，大幅增加GPU的利用率，从40%左右提升到了90%以上；
利用hfai高性能工具改造模型，将Alphafold2深度融合进幻方AI的训练平台中，大幅提高训练性能。

本次文章先和大家聊聊，如何优化蛋白质数据的处理，以获得高效的训练性能。

模型仓库：https://github.com/HFAiLab/alphafold-optimized

输入优化

Alphafold的输入特征可以分为这样三类：

蛋白质序列本身
Multi-Sequence Alignment(MSA)多序列对齐特征
Template蛋白质模板特征

它们在不同的蛋白质序列长度之间差异非常大，处理所需时间可以从几秒到上千秒不等，这会极大的影响模型训练的效率。例如在一次20分钟的训练过程中，单个长样本的处理就可能阻塞整体训练10几分钟。而在多机多卡大批次场景下，每个批次内有部分数据的处理时间极长几乎成了必然事件。

如上图所示，在训练Alphafold的初始阶段，输入的特征是通过序列数据库搜索，生物学工具（JackHmmer/HHSearch/…）计算而得出，需要依赖大量的CPU计算能力支持。在整个Alphafold模型训练的过程中，CPU算力往往成为了阻碍Alphafold训练性能提升的重要瓶颈。

针对上述问题，幻方AI采用特征预处理与特征裁切两种方案对输入进行优化。简单来说，特征预处理利用专门的CPU集群，将大量数据库搜索任务，生物学工具计算任务预先做掉，得到特征集和候选集；在模型训练阶段，模型直接拉取特征集和候选集，筛选特征，适当裁切长蛋白质序列特征。通过该方法我们实现了alphafold训练性能4-5倍的提升，在接下来的章节中我们会为大家详细描述。

1. 特征预处理

在一般的模型训练任务中，耗时的大头都在于GPU的前向推理和反向传播的计算，数据处理可以通过异步Dataloader载入的方式来掩盖时间开销。然而在Alphafold的训练中，模型对处理好的数据进行一次迭代需要大约11秒（基于A100 GPU），而对于这一个数据的预处理则需要花费大约150秒（一整颗EPYC 64核处理器），其是GPU运算时间的13倍，根本无法被GPU计算时所隐藏。在这种情况下，GPU利用率的理论上限也仅有7.3%，效率是非常低下的。同时，萤火二号集群上每个机器配备的是8张GPU和一张CPU，按照原有的多卡并行训练方式，CPU资源会更加紧缺，因此在训练时对特征处理进行优化已经毫无意义，我们需要寻找其他替代方案。

我们选择将特征处理单独分离出来，使用CPU集群算力预先完成一部分任务，将得到的候选样本数据进行保存。为了提高处效率，最大化CPU的利用率，我们使用Python多进程对特征预处理进行了并行加速。然而，特征处理涉及到多个不同的生物学工具，且每种工具对CPU的利用率各不相同，因此直接并行化很容易导致进程过多，从而产生大量的CPU调度抢占开销。根据统计，在对几个不同类型的特征进行处理时的CPU资源占用情况如下表所示：

蛋白质数据集	处理工具	耗时（秒）	CPU占用核数	总算力需求（分钟/核）
Uniref90	JackHmmer	567	4核	37.80
Magnify	JackHmmer	851	4核	56.73
Bfd-small	JackHmmer	308	4核	20.53
Pdb70	HHSearch	62	64核	66.13

从上表可以观察到：前三个数据集使用JackHmmer，其基本只能使用4核的CPU资源；而Pdb70使用HHSearch，其能够使用CPU上的所有资源。因此我们需要针对不同的处理工具建立不同的进程池进行调度分配。我们采用的优化策略是：同时发起多个JackHmmer任务，并每隔一段时间发起一个HHSearch任务。

在128核的CPU节点上，经过优化后的特征预处理的CPU占用率如下图所示：

可见此时已经可以基本以比较稳定的100%利用率使用CPU资源。

值得注意的是，特征预处理可以预先将蛋白质本身特征提取出来（序列信息、空间结构信息等），而对于MSA特征和Template特征，Alphafold的论文中提到，训练中是一个搜索过程，随机选取最合适的4个特征进行输入，因此预处理能提供适合适的候选集，而最终模型训练时选哪4个特征作为输入，预处理部分无法替代。

上述特征预处理出来的中间结果可以通过幻方AI自研的FFRecord格式进行存储，以充分利用萤火集群的高速文件系统3FS，详情可以参考之前的文章《幻方力量 | 高速文件系统：3FS》。

2. 特征裁切

前面提到，不同长度的蛋白质序列处理可能会导致训练被阻塞，除了将运行时处理转为预处理，我们还可以对较长的特征进行裁切。但需要注意的是，对特征进行特征裁切必然会影响模型训练的效果，因此需要在模型性能和训练效率二者间调节平衡。

特征裁切仅针对MSA特征和Template特征的处理，因为其主要目标是找到更多相似特征（序列、结构）以增强模型的表达能力。特征裁切可以有效减少对复杂蛋白质的特征处理开销，提高整个训练效率。蛋白质序列本身特征是不适用特征裁切的。

2.1 MSA特征裁切

MSA特征是一组与输入的人类蛋白质序列相似的其他生物蛋白质序列的集合，用于帮助模型Encoder部分形成蛋白质的特征向量。若输入的蛋白质序列长度为N个氨基酸，在预处理阶段中找到的同源蛋白质数量为M，则需要进行预处理的MSA特征的大小就为NxM。可见在蛋白质序列不能被裁切的前提下，只能通过对同源蛋白质数量进行裁切，即丢弃部分之前查找到的同源序列来减少这部分的耗时。

蛋白质序列长度N，MSA特征序列数M，最终处理时长T，他们满足如下线性关系： $T=C \times N \times M$

C是一个固定常数，和机器的处理能力有关。我们实验中发现，使训练不被阻塞的最大时长为150秒，因此我们可以根据上面的公式推得每个输入蛋白质序列的最大MSA特征数量。实际在实验中，只有不到10%的序列需要裁切MSA特征，因此这能够在尽可能不损害性能的情况下极大的提升训练的效率。

2.2 Template特征裁切

Template特征旨在抽取与要预测的蛋白质在空间结构上相似的蛋白质序列，用于帮助模型的Structure Module部分形成特征表达，输出蛋白质的空间结构。同样地，单个蛋白质可能有上百个模板，虽然在Alphafold的训练中最终只有4个模板需要被送入模型参与训练，但未必所有模板都与蛋白质真正匹配。因此一旦遇到大量模板都不符合要求的情况，在检查模板序列的过程中就可能要花费十几分钟的时间。

对于Template特征的处理可以进行并行加速，但受限于PyTorch Dataloader守护进程的问题和CPU资源量的限制，在实际训练中这种方法不太适用。因此我们同样考虑对Template特征进行裁切，以减少这步骤消耗的时间。

特征预处理获得Template特征候选集，按照匹配度进行排序；
训练时尽量使用与原蛋白质更匹配的模板；
对于模板匹配过程设有一个固定的时间限制，超时做裁切。

后面实验会展示上述方法在实际训练中的效果。

实验

首先，我们先测试Alphafold在单张A100显卡上进行训练时每一个Step的时间开销，作为基准：

Stage	Mean duration (s)	Percentage (%)
Training Step	12.65	100
Get Training Batch	1.87	14.78
Forward	4.24	33.52
Zero Grad	0.17	1.31
Exponential Moving Average	0.14	1.0
Optimizer Step	0.45	3.42
Backward	5.93	46.88

从表格中可以看到，训练一个Step就要花费12.65秒，其中最主要的时间开销在于Backward，Forward和Get Training Batch三步。由于使用的数据很少，并且没有使用多卡并行，此时特征处理带来的性能瓶颈和DDP的开销还没有能体现出来。

而当训练数据较多时，不同训练数据间的处理时间差异就会显得尤为突出。如下图所示，大多数输入的数据只需要数秒即可完成处理，然而少部分蛋白质在极端情况下会需要上百秒的预处理时间。

因此一旦遇到一条这样的输入数据，就会导致整个训练阻塞。而假设所有训练数据中只有5%会导致训练阻塞，那么在使用128的batch_size训练Alphafold（原论文参数）时，每个batch被阻塞的概率就为:

$P_{stuck} = 1-(1-0.05)^{128}=99.86\%$

这几乎在每个step的训练中都会遇到阻塞的情况。此时整个训练的GPU利用率会很极端的下滑到20%以下，其中大部分时间都花费在了等待batch中的某一个数据处理完成上。

例如在没有进行优化时使用16个节点，128张A100进行训练，那么前4个step的花费时间就为：413s, 177s, 100s, 101s, 每个step中的GPU利用率不到10%。可见此时若是不对数据处理进行优化，那么很难正常进行训练。

1. 特征预处理测试

我们在总共11万条的蛋白质序列训练数据中，随机抽取了3.5万条蛋白质序列进行解析，并统计了这些序列分别使用直接训练时处理和特征预处理方式的耗时：

蛋白质	原处理时间（s）	使用特征预处理后的时间（s）
3j3q	145.9	46.6
3j3y	129.8	41.0
6u42	78.9	27.5

如上表所示，对于处理时长大于10秒的序列，原来需要150秒的特征解析过程在优化后仅仅需要50秒左右，这大大抑制了超长序列对整个训练的阻塞影响。即使是小序列，也能得到将近3倍的速度提升。

接着，我们再看看使用特征预处理结果对训练时的GPU利用率和总训练时间能带来多大的收益。这里，我们先不进行特征裁切，先比较特征预处理方法后的性能提升。结果如下表所示：

	不进行特征预处理	进行特征预处理
每轮开始时的数据处理时长（s）	451	393
平均每个step时长（s）	50.3	23.57
训练被阻塞的次数	10	6
GPU利用率	19.58%	42.37%

可以看到，进行特征预处理能够获得非常明显的性能提升，GPU利用率提升了一倍。但同时，过多的训练阻塞导致整体的GPU利用率依然不够理想。

2. MSA特征裁切测试

接着，我们测试使用MSA特征裁切对训练时的GPU利用率和总训练时间的收益。这里MSA特征裁切时间常数为 2e-6，不进行Template特征裁切，都使用特征预处理。结果如下表所示：

	不进行MSA裁切	裁切MSA特征
每轮开始时的数据处理时长（s）	393	200
平均每个step时长（s）	23.57	10.87
训练被阻塞的次数	6	1
GPU利用率	42.37%	90.2%

可以看到，使用MSA特征裁切后GPU利用率能再提升一倍以上，训练阻塞也只发生一次。该方法能有效提升训练效率。

3. Template特征裁切测试

之后，我们再测试下Template特征裁切对训练效率的提升效果。时间限制为150秒，都使用特征预处理，不进行MSA特征处理。结果如下表所示：

	不进行Template裁切	裁切Template特征
每轮开始时的数据处理时长（s）	393	201
平均每个step时长（s）	23.57	11.65
训练被阻塞的次数	6	1
GPU利用率	42.37%	90.2%

使用Template特征裁切后GPU利用率同样能够提升一倍以上。

小结

作为2021年AI领域最具影响力的工作之一，Alphafold的复现无疑是有意义的，能帮助更多研究者和开发者们去探索这个方向的前景。通过优化数据输入，我们可以实现90%以上的GPU使用率，高效利用起了集群计算资源。同时，幻方AI也很多高效的训练工具，比如hfreduce，3FS，算子等等，它们是否能使得Alphafold的训练进一步加速？下一期技术分享将为大家呈现。

hfai.nn | 性能卓越的深度学习算子

High-Flyer — Thu, 23 Dec 2021 21:00:00 GMT

深度学习框架的流行（如PyTorch，Tensorflow等）极大方便了我们研发设计各种各样的AI模型，而在实际落地的环节中，孵化于实验室里的模型代码往往在生产环境上面临着性能、准确度、资源等各种各样的问题。随着AI与现实世界的业务结合愈加密切，高性能且易使用的深度学习算子愈发受到AI研究者和开发者们的青睐。

幻方AI依托萤火二号集群，自研了一系列AI基础设施：3FS，hfreduce，ffrecord等。如之前系列文章中介绍的，他们极大地提升了模型训练中数据读取的性能，充分利用GPU的算力。而对于计算环节，幻方AI同样对Pytorch框架进行了深度优化，结合萤火二号的集群特点，对一些常用的AI算子重新研发，提升效率，进一步提升了模型整体的训练效率。

本期文章将为大家介绍，幻方AI的深度学习算子，hfai.nn，供大家使用。

hfai.nn.parallel

在之前的文章《幻方萤火 | 模型并行训练工具 hfreduce》中介绍了幻方AI自研的并行训练工具，设计的全新的通讯模式替代NCCL，在幻方萤火集群上取得非常好的加速效果。这里，我们对接口按照PyTorch框架风格进一步封装，降低大家的使用门槛。

hfai DDP

我们提供了 hfai.nn.parallel.DistributedDataParallel 接口，其在 hfreduce 的基础上进行封装，使用方法和 PyTorch 的 DDP 基本相同。在萤火二号集群的环境下，hfai DDP 多机并行的性能比 PyTorch DDP 更好。

使用案例：

from hfai.nn.parallel import DistributedDataParallel

# init dist
ip = os.environ["MASTER_ADDR"]
port = os.environ["MASTER_PORT"]
hosts = int(os.environ["WORLD_SIZE"])  # number of nodes
rank = int(os.environ["RANK"])  # node id
gpus = torch.cuda.device_count()  # gpus per node

dist.init_process_group(
    backend="nccl",
    init_method=f"tcp://{ip}:{port}",
    world_size=hosts * gpus,
    rank=rank * gpus + local_rank
)
torch.cuda.set_device(local_rank)

model = MyModel().cuda()
model = DistributedDataParallel(model, device_ids=[local_rank])

hfai.nn.to_hfai/to_torch

有了上述分布式训练管理工具，接下来就是具体的优化算子了。为了让模型计算的更快，我们优化CUDA kernel，融合多运算，优化模型结构，使得深度学习算子计算更加的高效。

具体优化的算子包括：

torch.nn.lstm - hfai.nn.lstm
torch.nn.MultiheadAttention - hfai.nn.MultiheadAttention
torch.nn.LayerNorm- hfai.nn.LayerNorm

使用方法非常简单，幻方AI提供了工具可以一行代码进行算子的切换，即to_hfai和to_torch。

to_hfai

to_hfai是方便一键使用hfai算子的工具，会对你的model进行一次深度遍历，将其中能做转化Module替换成hfai的。转化后的算子在training以及inference的时候都会有可观的性能提升。

to_torch

to_torch是将使用hfai算子的model转化回使用torch算子的工具，以便在非 hfai 下使用训练出来的模型。设计思路同to_hfai。

使用案例：

from hfai.nn import to_hfai,to_torch

# get model from hfai && to_hfai
model = hfai.models.bert_large_MLM(vocab_size=vocab_size, max_length=max_length).cuda()
model = to_hfai(model)

性能测试

那么，幻方AI优化的算子具体能达到多少的加速效果呢？这里，我们把性能测试数据给出：

LSTM

参数配置：

input_size: 985
hidden_size: 985
bias: True
num_layers: 1
batch_first: True
dropout: 0
bidirectional: False
batch_size: 4
seq_len: 1000

连续测试10轮，统计整体每轮的耗时，测试结果如下：

torch.nn.LSTM
- forward: 31.78 ms/epoch
- backard: 93.60 ms/epoch
- all: 125.38 ms/epoch
hfai.nn.LSTM
- forward: 4.82 ms/epoch
- backard: 9.25 ms/epoch
- all: 14.07 ms/epoch

forward 提升6倍，backward 提升了10倍，整体训练提升了8倍！

Attention

参数配置：

batch_size: 200
seq_len: 256
embed_dim: 1024
attn_heads: 16

连续测试10轮，统计forward，backward，还有inference的耗时，测试结果如下：

torch.nn.MultiheadAttention
- forward: 17.581 ms/epoch
- backard: 21.986 ms/epoch
- all: 39.567 ms/epoch
- inference: 17.372 ms
hfai.nn.MultiheadAttention
- forward: 12.165 ms/epoch
- backard: 16.515 ms/epoch
- all: 28.680 ms/epoch
- inference: 12.112 ms

forward 提升44%，backward 提升了33%，整体训练提升了37%，inference也提升43%！

同时，当batch_size比较大，seq_len比较小，以及输入qkv一样的情况下，hfai.nn.MultiheadAttention 同样具备优异性能。

参数配置：

batch_size: 5000
seq_len: 64
embed_dim: 512
attn_heads: 8

连续测试10轮，统计forward，backward，还有inference的耗时，测试结果如下：

torch.nn.MultiheadAttention
- forward: 22.713 ms/epoch
- backard: 41.831 ms/epoch
- all: 64.543 ms/epoch
- inference: 22.925 ms
hfai.nn.MultiheadAttention
- forward: 19.599 ms/epoch
- backard: 24.379 ms/epoch
- all: 43.978 ms/epoch
- inference: 19.037 ms

forward 提升15%，backward 提升了71%，整体训练提升了46%，inference也提升20%！

LayerNorm

参数配置

batch_size = 200
seq_len = 128
hidden_size = 1024

连续测试10轮，统计forward，backward，还有inference的耗时，测试结果如下：

torch.nn.LayerNorm
- forward: 0.514 ms/epoch
- backard: 0.217 ms/epoch
- all: 2.467 ms/epoch
hfai.nn.LayerNorm
- forward: 0.217 ms/epoch
- backard: 0.121 ms/epoch
- all: 1.306 ms/epoch

forward 提升137%，backward 提升了79%，整体训练提升了88%！

总结

通过优化CUDA kernel和模型结构，融合多运算操作，幻方AI优化的深度学习算子在实际测试中表现出惊人的加速效果，优化的算子与幻方AI基础设施结合，让模型训练高效迅速，计算资源充分利用，加速更多场景的探索与研究。算子优化一直持续中，非常欢迎更多优秀同行加入我们一同共建。

萤火跑模型 | Alphafold 蛋白质结构预测

Vachel — Fri, 05 Nov 2021 21:00:00 GMT

2021年刚刚走过，要说去年一年AI领域的最重大进展，DeepMind在Nature上发表的“AlphaFold”无出其右，其旨在根据一个蛋白质的氨基酸序列来确定它的3D结构，是生物学领域中的几个重大挑战课题之一。在“蛋白质结构预测关键评估”（CASP）的蛋白质结构预测年赛上，AlphaFold击败了其他百来支团队并取得了惊人的效果，有些时候，AlphaFold预测的结构与利用X射线晶体学和近几年的冷冻电镜（cryo-EM）等“金标准”实验方法所确定的结构几乎一样。科学家说，目前看来，AlphaFold暂时还不能取代这些费力又昂贵的实验方法，但它将带来全新的研究生命的方式。

作为一个计算机专业的算法从业人员，本期笔者带着浓厚的兴趣来尝试一下 DeepMind 开源的 Alphafold 模型并验证效果，试跑的平台是幻方AI团队研发的“萤火二号”。我们复现了 Alphafold 的训练与预测过程，为大家带来第一手的测试体验。

论文标题：Improved protein structure prediction using potentials from deep learning

原文地址：https://www.nature.com/articles/s41586-019-1923-7

源码地址：https://github.com/aqlaboratory/openfold

模型仓库：https://github.com/HFAiLab/alphafold-optimized

模型介绍

AlphaFold 解决的问题是蛋白质折叠问题。输入是一个氨基酸序列，每一个位置代表一个元素，输出是一个拓扑结构，如下图所示：

直接输出一个拓扑结构对于深度学习来说比较困难，一般需要输出结构化的数据。DeepMind 将该问题转化为预测氨基酸链的一些性质。准确的来说，深度模型输出的是每一个氨基酸单元和其下一个氨基酸单元在空间中的夹角与距离，再组合成拓扑结构。

如上图所示，AlphaFold算法大致分为以下几个部分：

特征工程：序列和MSA特征抽取，结合专家经验数据库，把氨基酸链的输入转换到特征空间；
深度神经网络结构预测：依据1中的特征预测氨基酸链的一些性质，比如氨基酸之间两两的距离分布，氨基酸链的夹角分布；
Potential Construction：结合专家经验构造一个评估函数，来评估2中神经网络输出解的合理程度；
结构生成：对于2中预测的距离分布、夹角分布，使用3中的评估函数评估Loss，然后使用梯度下降法优化，直到收敛。

可以看到，Alphafold 核心贡献在于：

深度学习模型融入专家经验：特征构建阶段，尽可能结合领域知识构建特征，以提高模型表现；评估阶段，将专家知识（氨基酸链的物理特性）结合到损失函数中，进行更精准的判断。
全新的蛋白质折叠问题的建模方式：DeepMind不是直接去预测蛋白质的结构，而是将该问题转化为预测氨基酸折叠后的一些性质，再将这些性质转化为可微的约束，并求解。这种思路可以应用在很多复杂的求解问题上。
对于梯度下降算法的优化：DeepMind提出通过类似进化/退火算法，解决梯度下降中的局部最优问题，以求解真正的“最优解”。

模型实践

DeepMind 仅开源了 Alphafold 的预测代码和预训练好的模型参数，而对于训练部分，OpenFold 与 UniFold 分别开源了 Pytorch 和 Tensorflow 的版本。我们选择 OpenFold 进行改造，因为其复现完整度比较高，且使用 Pytorch 框架更适配萤火集群的运行环境。项目地址： https://github.com/aqlaboratory/openfold

1. 数据集

从上文可以知道，Alphafold 数据集包含两个部分：氨基酸序列数据、专家经验数据。具体的，数据说明如下：

bfd / small_bfd：Big Fantastic Database，由论文作者 Martin Steinegger 构建的蛋白质序列数据库，包含65,983,866个蛋白质家族（MSAs）和2,204,359,010个蛋白质序列（HMMs），总计7个文件1.8T数据量。
mgnify：一个微生物基因组数据库，使用JackerHammer工具进行检索，总计一个文件64G数据量。
pdb70：一个蛋白质结构数据库，总计9个文件56G数据量。
pdb_mmcif：一个蛋白质结构数据库，包含超过55,000种不同蛋白质的180,000多个3D结构，总计约18万个文件206G数据量。
uniclust30：1.3亿条氨基酸序列，使用hhsuite中的HHblits工具进行快速msa检索，总计13个文件86G数据量。
uniref90：1.3亿条氨基酸序列，是msa检索所需要的库，使用JackerHammer工具进行检索，总计1个文件58G数据量。

在萤火集群上，这些数据如果以独立文件形式直接用于训练，小文件的高频读写会给集群的存储系统造成不小压力。对应的，我们需要将数据转化成 FFRecord 数据格式，合并多个小文件以减小磁盘 IO 的开销。FFRecord (FireFlyer Record) 数据格式是幻方AI自研的简单高效的存储二进制记录的文件格式。为了方便研究者的使用，我们将以上的多个数据集进行整合，加入hfai 高性能数据集仓库，对外提供 hfai.datasets.AlphafoldData 接口。具体的使用案例如下：

from hfai.datasets import AlphafoldData

# 将数据转为Alphafold训练的格式
transform = AlphafoldDataTransform()
dataset = AlphafoldData(transform)
dataloader = dataset.loader(batch_size=1, num_worker=8)

for data in dataloader:
    # training iter

AlphafoldData 数据集接口返回的数据是一个元组，其中包含了蛋白质本身的信息 pdb_code， mmcif_string 和在蛋白质数据库中匹配得到的结果 bfd_hits，mgnify_hits，pdb70_hits，uniref90_hits。这些数据在训练中需要被 AlphafoldDataTransform 进一步处理，具体的使用案例可以访问我们开源在模型仓库中的 Alphafold-optimzied。

除了以上整合的数据，pdb_mmcif 数据在 Alphafold 训练后期指导着模型生成合适的蛋白质结构，其具有将近 18w 个小文件，同样需要将其转化成 FFRecord 数据格式。

原始数据转 FFRecord：

from ffrecord import FileWriter

out_file = os.path.join(mmcif_path, "mmcifs.ffr")
writer = FileWriter(outfile, 184929)
num, mapids = 0, {}
for file in os.listdir(os.path.join(mmcif_path, "mmcif_files")):
    with open(os.path.join(mmcif_path, "mmcif_files", file), "r") as fp:
        mmcif_string = fp.read()
        ms = pickle.dumps(mmcif_string)
        writer.write_one(ms)
        mapids[file[:-4]] = num
        num += 1
with open(os.path.join(mmcif_path, "mmcifs_index.pk"), "wb") as fp:
    pickle.dump(mapids, fp)

读取 FFRecord 文件：

from ffrecord import FileReader

mmcif_reader = FileReader(os.path.join(mmcif_dir, "mmcifs.ffr"), check_data=True)
with open(os.path.join(mmcif_path, "mmcifs_index.pk"), "rb") as fp:
    mmcif_index = pickle.load(fp)
if hit_pdb_code in mmcif_index:
    idx = [mmcif_index[hit_pdb_code]]
    data = mmcif_reader.read(idx)
    cif_string = pickle.loads(data[0])

以上数据皆已集成进hfai 数据集仓库，如果需要利用该数据集进行科研的朋友可以直接拉取。

2. 预处理

从上文我们知道，Alphafold 需要与专家经验库进行交互，抽取 MSA 特征。这里用到了多个生物学工具：

openmm：用于分析模拟的高性能工具包，是 Omnia（预测生物分子模拟的工具套件）的一部分。
kalign：一款针对大规模基因组序列的多序列对比工具，运行速度与准确度都比较高。
jackhmmer：一款高效的蛋白质搜索工具，用于从序列数据库中，找到同源的序列。
pdbfixer：用于修复PDB数据，对其中蛋白质结构进行基本的准备或者修复。
HHsuite：一款蛋白质研究软件，快速搜索分析蛋白质的功能特性。

以上工具都只能使用CPU多核进行加速计算，对CPU的算力比较依赖。对应的，在幻方萤火的系统结构中，GPU算力池与CPU算力池是分离开的。因此，数据预处理与模型计算要拆分成两个模块单独进行。

name: precompute_alignments
code:
	directory: /ceph-jd/prod/jupyter/hwj/notebooks/openfold
  entrypoint: scripts/precompute_alignmnets.py
  parameters: mmcif/ alignments/ --cpus 250
resource:
	- container: cuda_11
  - node_count: 1
    group: jd_dev_alpha
priority: 20

这里，jd_dev_alpha 是CPU计算池。通过 hfai run提交上述yaml 文件便可以发起对于目标蛋白质结构的预处理任务。

3. 模型训练

openfold 采用了 Pytorch_lightning 来构建整个训练代码，其自动封装了 dataloader、多机多卡训练、checkpoint 保存与打断自恢复等诸多逻辑，可以直接适配萤火二号的智能分时调度系统。这里我们通过 hfai run运行下述Yaml文件：

name: precompute_alignments
code:
	directory: /ceph-jd/prod/jupyter/hwj/notebooks/openfold
  entrypoint: train_openfold.py
  parameters: mmcif/ alignments/ data/pdb_mmcif output 2021-10-10 --gpus 1 --resume_from_ckpt output/resume_checkpoint.ckpt
resource:
	- container: cuda_11
  - node_count: 1
    group: jd_a100
priority: 20

这里，jd_a100是GPU计算池。由于目标预测的样本不多，这里使用一个计算卡进行训练，过程如下图所示：

训练结束后，我们加载输出的训练模型，预测蛋白质的结构，得到PDB文件。通过PyMol3d 工具，可以将PDB中的蛋白质结构数据可视化出来，最终效果如下：

体验总结

Alphafold作为今年火爆全球的人工智能领域重点突破之一，开启了AI+生物医药领域的新时代。借助萤火二号和自研文件系统，我们可以比较高效地训练该模型。同时，该模型不同于以往深度学习训练场景，其特征抽取与训练，对CPU与GPU的计算需求都比较高。在我们的测试中，GPU整体使用率不高，在20%~30%左右，训练的时间瓶颈是在CPU的计算上。我们欢迎感兴趣的学者们加入我们，继续在幻方AI团队的萤火平台上对该课题进行进一步探索。综合体验打分如下：

研究指数：★★★★★

该模型开创了AI+生物的新课题，引领了一个新的AI时代。
开源指数：★★★

数据已开源，代码开源了预测部分，训练代码开源需要依赖第三方实现。
门槛指数：★★★★★

数据量大，需要CPU与GPU双重算力支持，对算力的需求比较大。
通用指数：★★

模型专注于蛋白质结构预测，有大量的专家经验结合，通用性比较差。
适配指数：★★

需要CPU与GPU的双重算力，整体GPU利用率不高。同时代码和数据都有一定改造工作。

幻方 AI 紧跟 AI 研究的前沿浪潮，致力于用领先算力助力AI落地与价值创造，欢迎各方数据研究者与开发者们一同共建。

萤火跑模型 | Autoformer 长时序预测

Vachel — Mon, 11 Oct 2021 15:00:00 GMT

长时间序列预测（Long-Term Series Forecasting，以下称为 LTSF）在现实世界中有非常广阔的应用场景，例如电力资源预估、疾病传播与扩散、经济发展预测等实际问题。前不久分享的《萤火跑模型 | Informer 上手实践》介绍了获得AAAI 2021最佳论文奖的Informer模型，以及其在幻方萤火二号上的试跑和优化。该工作改造Transformer模型来解决LTSF领域下的问题，取得优异的效果，带动起了该方向的研究热情。不久之后的NeurIPS 2021，来自清华的几位学者结合之前和工作和Informer的模型设计，提出了Autoformer模型，在相同预测任务上大幅超越了之前的模型，实现了 38% 的相对效果提升。

幻方AI最近对这项工作进行了整理和优化，在幻方萤火二号上复现了Autoformer的实验。通过幻方自研的3FS、hfreduce、算子等，对模型训练进行提速，整合代码简化接口，融入hfai数据仓库和模型仓库中，帮助研究者和开发者们降低使用门槛。本期文章将为大家详细描述。

论文标题：Autoformer: Decomposition Transformers with Auto-Correlation for Long-Term Series Forecasting

原文地址：https://arxiv.org/abs/2106.13008

源码地址：https://github.com/thuml/Autoformer

模型仓库：https://github.com/HFAiLab/ltsf-former

模型介绍

我们知道，Transformer是通过自注意力机制（self-attention）来捕捉时刻间的依赖，在NLP领域上取得了显著效果。但是在长期序列预测中，仍存在不足：

长序列中的复杂时间模式使得注意力机制难以发现可靠的时序依赖；
基于Transformer的模型不得不使用稀疏形式的注意力机制来应对二次复杂度的问题，但造成了信息利用的瓶颈。

为突破上述问题，Autoformer主要做了以下创新：

突破将序列分解作为预处理的传统方法，提出深度分解架构（Decomposition Architecture），能够从复杂时间模式中分解出可预测性更强的组分；
基于随机过程理论，提出自相关机制（Auto-Correlation），代替点向连接的注意力机制，实现序列级连接和较低复杂度，打破信息利用瓶颈。

相比与Informer，他们都提出了要改变注意力机制，提升序列级连接和降低运算复杂度。不同的是，Autoformer进一步提出，将时序进行分解，从复杂模式中抽取出可预测性更强的成分。具体是如何做的呢？下面的章节会为大家简单介绍。

深度分解架构

时间序列分解是指将时间序列分解为几个组分，每个组分表示一类潜在的时间模式，如周期项（seasonal），趋势项（trend-cyclical）。由于预测问题中未来的不可知性，通常先对过去序列进行分解，再分别预测。但这会造成预测结果受限于分解效果，并且忽视了未来各个组分之间的相互作用。

Autoformer是将序列分解作为一个内部单元，嵌入到编码-解码器中。在预测过程中，模型交替进行预测结果优化和序列分解，即从隐变量中逐步分离趋势项与周期项，实现渐进式分解。在编码器中，基于滑动平均思想逐步分解出各成分，在解码器中，对各个部分分别建模。基于这种渐进式分解架构，模型可以在预测过程中逐步分解隐变量，并通过自相关机制、累积的方式分别得到周期、趋势组分的预测结果，实现分解、预测结果优化的交替进行、相互促进。

自相关机制

Autoformer的自相关机制用来实现高效的序列级连接，从而扩展信息效用。一般不同周期的相似相位之间通常表现出相似的子过程，Autoformer利用这种序列固有的周期性来设计自相关机制，其中，包含基于周期的依赖发现（Period-based dependencies）和时延信息聚合（Time delay aggregation）。

基于周期的依赖发现采用的是随机过程理论，计算相关性作为未归一化的周期估计的置信度。而时延信息聚合是将相似的子序列信息根据计算出的周期长度进行聚合，以实现序列级的连接。

Autoformer的作者们在文中给出了使用该方法后的效果。如上图所示，相比于之前的注意力机制或者稀疏注意力机制，自注意力机制（Auto-Correlation）实现了序列级的高效连接，从而可以更好的进行信息聚合，打破了信息利用瓶颈。

模型实践

Autoformer的作者对模型进行了完整的开源，幻方AI的研究者们对其进行了试跑复现和优化。

数据集

Informer和Autoformer基于相同的数据集进行试验，包括电力，经济，交通，医疗等若干个长时序数据集。历史分析长度，预测长度，预测范围作为参数可以动态调整。这里，幻方AI将这些能通用的数据集进行整合，将原始数据转化成FFRecord格式，用户调用简单的接口便可以加载高性能的dataloader进行模型训练。具体的，如下例所示：

from hfai.datasets import LTSF

dataset = LTSF(data_name='ETTh1', split='train', seq_len=96, label_len=48, pred_len=24, features='S', target='OT', timeenc=0, freq='h')
loader = dataset.loader(batch_size=64, num_workers=4)

for x, y, x_mark, y_mark in loader:
    # training model ...

如上例代码，我们加载hfai数据仓库中的ETTh1.csv数据，hfai数据仓库会自动切分好数据集，进行特征工程，并返回高性能的dataloader。

模型训练

幻方AI对作者开源的模型源码进行了优化，以hfreduce代替NCCL，以自研的高性能算子替代模型中的LayerNorm，Attention等算子，实现Autoformer在萤火二号上的高性能并行训练。用户可以直接在hfai模型仓库中直接调用到我们优化后的模型：

from hfai.models import Autoformer
import hfreduce.torch as hfr

model = Autoformer(enc_in=1, dec_in=1, c_out=1, seq_len=96, label_len=48, out_len=24, e_layers=2, d_layers=1, dropout=0.05, embed='timeF').cuda()
loss = Loss(model)
optimizer = Optimizer(model)

reducer = hfr.AsyncReduceFloat(registry_host, port, proc_rank, procs_per_node, node_rank, node_cnt, model)

# each epoch
for x, y, x_mark, y_mark in loader:
    reducer.zero_grad() 
    
    model(x, y, x_mark, y_mark)
    loss.backward()
    
    reducer.synchronize()

由于开源数据不算很大，我们试跑时只申请一个节点，8张A100显卡并行训练。通过优化后训练效果很高，很快便可以完成试验：

可以看到，每Epoch耗时仅2s左右，相比单卡性能提升了8倍以上。GPU利用率在95%左右，充分利用其了A100的计算性能。模型在第16轮基本达到了收敛状态。

体验总结

Autoformer进行了完整的开源，代码结构清晰，比较容易与幻方的训练优化工具结合，获得效果和性能的双提升。同时，该模型承接Informer，继续推动了“transformer+时序预测”这个方向的研究走向深入，助力更多真实场景的提效与进步。

综合体验打分如下：

研究指数：★★

该模型着眼于长时间序列的预测问题，是一个广泛研究的基础课题。
开源指数：★★★★★

数据和代码都已开源，代码逻辑清晰可读性高。
门槛指数：★★★

数据量小，模型计算、内存和体系结构均已优化，普通高性能PC即可运行。
通用指数：★★★★

模型是对Transformer的优化，将时序分解融入，在很多场景下验证了其效果。
适配指数：★★★★★

依赖简单，很容易与幻方AI的训练优化工具结合，提效明显。

幻方 AI 紧跟 AI 研究的前沿浪潮，致力于用领先算力助力AI落地与价值创造，欢迎各方数据研究者与开发者们一同共建。

PyTorch 分布式训练方法

Vachel — Fri, 18 Jun 2021 21:00:00 GMT

2018年，将近3亿参数的Bert模型横空出世，将NLP领域推向了新的高度。近年来人工智能领域的发展愈来愈趋向于对大模型的研究，各大AI巨头都纷纷发布了其拥有数千亿参数的大模型，诞生出了很多新的AI应用场景。另一方面，多种因素继续推动大模型的长足发展：1) 社会正经历着深度的数字化转型，大量的数据逐渐融合，催生出许多人工智能的应用场景和需求；2) 硬件技术不断进步：英伟达 A100 GPU，Google的 TPU，阿里的含光800等，推动着AI算力的不断提升。可以这么说，“大数据+大模型+大算力”的组合，是迈入 AI 2.0 的基石。然而，对于大多数AI从业者们来说，受限于一些因素，平常很难接触到大数据或者大算力，很多高校研究生们往往都习惯于用一张显卡（个人PC）来加速训练，这不仅训练缓慢，甚至会因此限制住了想象，禁锢了创造力。幻方AI根据自身业务需求，构建了“萤火二号”，其旨在提供澎湃如海的深度学习算力服务。团队自研了智能分时调度系统、高效存储系统与网络通信系统，可以将集群当成一台普通计算机来使用，根据任务需求弹性扩展GPU算力；自研了hfai数据仓库、模型仓库，优化AI框架与算子，集成落地了许多前沿的应用场景；通过Client接口或者Jupyter就可以轻松接入，加速训练起AI大模型。

本期文章分享的，是如何使用起多张显卡，来加速你的AI模型。分布式训练技术逐渐成为AI从业者必备技能之一，这是从“小模型”走向“大模型”的必由之路。我们以 PyTorch 编写的ResNet训练为例，为大家展示不同的分布式训练方法及其效果。

训练任务：ResNet + ImageNet

训练框架：PyTorch 1.8.0

训练平台：幻方萤火二号

训练代码：https://github.com/HFAiLab/pytorch_distributed

训练准备

萤火二号当前完整支持了PyTorch的并行训练环境。笔者先使用1个节点8张卡，测试不同的分布式训练方法所需要的时间、显卡利用效率；之后采用多个节点，来验证并行加速的效果。测试的方法如下：

nn.DataParallel
torch.distributed + torch.multiprocessing
apex 半精度

为充分利用显存，这里batch_size设为400，记录每Epoch的耗时来进行对比。我们先采用单机单卡测试作为baseline，每Epoch耗时1786秒上下。

测试结果发现，有半精度加持的并行加速效果最好；DataParallel较慢，且显卡资源利用度不高，不推荐使用；多机多卡带来的加速效果非常显著。整体结果如下：

nn.DataParallel

DataParallel 是PyTorch早期提出的数据并行训练方法，其使用单进程控制，将模型和数据加载到多个GPU中，控制数据在 GPU 之间的流动，协同不同 GPU 上的模型进行并行训练。

使用起来非常方便，我们只需要用 nn.DataParallel 包装模型，再设置一些参数即可。需要定义的参数包括：

参与训练的 GPU 有哪些，device_ids=gpus；
用于汇总梯度的 GPU 是哪个，output_device=gpus[0] 。

DataParallel 会自动帮我们将数据切分 load 到相应 GPU，将模型复制到相应 GPU，进行正向传播计算梯度并汇总：

值得注意的是，这里因为有8张显卡，**batch_size**要设为3200，由DataParallel均分到不同的显卡里。模型和数据都要先 load 进 GPU中，DataParallel 的 module 才能对其进行处理，否则会报错。

发起训练。从下图我们可以看到，在显存几乎占满的情况下，8张显卡平均GPU利用率不高（36.5%）。

单独查看每一张显卡的利用情况，可以发现资源使用不平均。其中0号显卡作为梯度汇总，资源利用度相比其他显卡要高一点，整体利用率不到50%。

最终，nn.DataParrallel方法下，ResNet平均每Epoch耗时984s左右，相比单机单卡加速大约两倍效果。

torch.distributed

在 pytorch 1.0 之后，PyTorch官方对分布式的常用方法进行了封装，支持 all-reduce，broadcast，send 和 receive 等等。通过 MPI 实现 CPU 通信，通过 NCCL 实现 GPU 通信，解决了 DataParallel 速度慢，GPU 负载不均衡的问题。

与 DataParallel 的单进程控制多 GPU 不同，在 distributed 的帮助下，我们只需要编写一份代码，torch 就会自动将其分配给 n 个进程，分别在 n 个 GPU 上运行。

具体的，首先使用 init_process_group 设置GPU之间通信使用的后端和端口：

然后，使用 DistributedDataParallel 包装模型，它能帮助我们为不同 GPU 上求得的梯度进行 all reduce（即汇总不同 GPU 计算所得的梯度，并同步计算结果）。all reduce 后不同 GPU 中模型的梯度均为 all reduce 之前各 GPU 梯度的均值：

model = torch.nn.parallel.DistributedDataParallel(model.cuda(), device_ids=[local_rank])

最后，因为是采用了多进程管理，需要使用 DistributedSampler 对数据集进行划分，它能帮助我们将每个 batch 划分成几个 partition，在当前进程中只需要获取和 rank 对应的那个 partition 进行训练：

值得注意的是**，与****nn.DataParrallel**不同，这里**batch_size**应该为400，因为其只负责当前rank下对应的partition，8张卡一共组成3200个样本的batch_size。

至于 API 层面，PyTorch 为我们提供了 torch.distributed.launch 启动器，用于在命令行分布式地执行 python 文件。在执行过程中，启动器会将当前进程的（其实就是 GPU的）index 通过参数传递给 python，我们可以这样获得当前进程的 index：

parser = argparse.ArgumentParser()
parser.add_argument('--local_rank', default=-1, type=int, help='node rank for distributed training')
args = parser.parse_args()
print(args.local_rank)

启动8个命令行，执行如下命令：

CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 python -m torch.distributed.launch --nproc_per_node=8 train.py

torch.multiprocessing

手动启动命令行比较麻烦，这里有一种更简单的方法：torch.multiprocessing 会帮助自动创建进程，绕开 torch.distributed.launch 自动控制开启和退出进程的一些小毛病。

如下面的代码所示，将原本需要 torch.distributed.launch 管理的执行内容封装进 main 函数中，通过 torch.multiprocessing.spawn 开启了 nprocs=8 个进程，每个进程执行 main 并向其中传入 local_rank（当前进程 index）。

发起训练。从下图我们可以看到，显存几乎占满，8张显卡平均利用率高（95.8%）。

同时，每张显卡利用都比较充分，利用率都在80%以上：

最终，torch.distributed方法下，ResNet平均每Epoch耗时239s左右，相比单机单卡加速大约八倍效果。

apex

Apex 是 NVIDIA 开源的用于混合精度训练和分布式训练库。Apex 对混合精度训练的过程进行了封装，改两三行配置就可以进行混合精度的训练，从而大幅度降低显存占用，节约运算时间。此外，Apex 也提供了对分布式训练的封装，针对 NVIDIA 的 NCCL 通信库进行了优化。PyTorch 1.6版本之后，开始原生支持，即torch.cuda.amp。

直接使用 amp.initialize 包装模型和优化器，apex 就会自动帮助我们管理模型参数和优化器的精度了，根据精度需求不同可以传入其他配置参数。

在分布式训练的封装上，Apex 的改动并不大，主要是优化了 NCCL 的通信。因此，大部分代码仍与 torch.distributed 保持一致。使用的时候只需要将 torch.nn.parallel.DistributedDataParallel 替换为 apex.parallel.DistributedDataParallel 用于包装模型。

在正向传播计算 loss 时，Apex 需要使用 amp.scale_loss 包装，用于根据 loss 值自动对精度进行缩放：

发起训练。从下图我们可以看到，在同等batch_szie条件下，显存只占用了60%，而8张显卡平均利用率高（95.8%）。

同时，每张显卡利用都比较充分，利用率都在80%以上：

最终，apex方法下，ResNet平均每Epoch耗时230s左右，相比torch.distributed加速了一点。同时，apex也节约了GPU的算力资源，对应的可以设置更大的batch_size，获取更快的速度。我们极限测试下，能在torch.distributed基础上再提升50%的性能左右。

多机多卡训练加速

从上面的试验可以看到，8张卡带来的加速效果近似8倍。如果我们使用更多的机器，是否会有更快的提速效果呢？这里，我们对 apex的代码，配置4个节点，一共32张A100 GPU来进行试验，获得了惊人的效果。可以看到，每Epoch总耗时52秒左右，相比单机单卡速度提升30多倍，几乎和显卡数量保持一致。

分布式评价

以上展示了分布式训练的过程。然而，如何对训练的结果进行快速的推理与测试，也是非常重要的问题，例如：

训练样本被切分成了若干个部分，被若干个进程分别控制运行在若干个 GPU 上，如何在进程间进行通信汇总这些（GPU 上的）信息？
使用一张卡进行推理、测试太慢了，如何使用 Distributed 进行分布式地推理和测试，并将结果汇总在一起？
…

要解决这些问题，我们需要一个更为基础的 API，汇总记录不同 GPU 上生成的准确率、损失函数等指标信息。这个 API 就是 torch.distributed.reduce。

reduce等 API 是 torch.distributed 中更为基础和底层的 API。这些 API 可以帮助我们控制进程之间的交互，控制 GPU 数据的传输。在自定义 GPU 协作逻辑，汇总 GPU 间少量的统计信息时，大有用处。熟练掌握这些 API 也可以帮助我们自己设计、优化分布式训练、测试流程。

如上图所示，它的工作过程包含以下两步：

在调用 reduce(tensor, op=…)后，当前进程接受其他进程发来的 tensor。
在全部接收完成后，当前进程（例如rank 0）会对当前进程的和接收到的 tensor 进行 op 操作。

通过上述步骤，我们就能够对不同 GPU 上的训练数据的损失函数进行求和了：

除了reduce，PyTorch官网提供了诸如 scatter, gather, all-reduce等6种集合通信方案，具体的参考官方文档： https://ptorch.com/docs/1/distributed

总结

本文介绍了不同的PyTorch并行训练方法，并在幻方萤火二号上进行了测试。从测试结果中我们能够得出，多机多卡并行训练能有效提升我们效率，apex 方法结合torch.distributed 能最大限度的发挥显卡的算力，让加速效果与显卡数量一致。同时，我们也发现，随着显卡数量的增多，受限于梯度汇合中的reduce计算，单个显卡的利用率会逐渐降低。如何高效优化reduce算法，尽可能提高gpu算力，从而加速整体训练效率，是接下来值得大家继续研究的课题。

萤火跑模型 | Informer 上手实践

Vachel — Sun, 18 Apr 2021 21:00:00 GMT

长时间序列预测（Long Sequence Time-Series Forecasting，以下称为 LTSF）在现实世界中是比较基础但又十分重要的研究场景，例如商品销量及库存的预测、电力消耗的规划、股票价格、疾病传播与扩散等实际问题。然而LTSF因为其历史数据量大、计算复杂性高、预测精度要求高，一直以来并没有取得太好的效果。

今年人工智能顶级大会AAAI的最佳论文奖项中有一篇来自北京航空航天大学的工作：Informer，其主要的工作是改造 Transfomer 算法来实现LTSF，并开源了代码与数据。笔者最近在幻方AI的萤火平台上尝试复现了该论文的实验，为大家带来第一手的测试体验。

论文标题：Informer: Beyond Efficient Transformer for Long Sequence Time-Series Forecasting

原文地址：https://ojs.aaai.org/index.php/AAAI/article/view/17325

源码地址：https://github.com/zhouhaoyi/Informer2020

模型仓库：https://github.com/HFAiLab/ltsf-former

模型介绍

近年来的研究表明，Transformer具有提高预测能力的潜力。然而，Transformer也存在几个问题，使其不能直接适用于LTSF问题，例如时间复杂度、高内存使用和“编码-解码”体系结构的固有局限性。为了解决这些问题，作者基于Transformer设计了一种适用于LTSF问题的模型，即Informer模型，该模型具有三个显著特征:

ProbSpare self-attention机制，有效降低了时间复杂度和内存使用量。
通过将级联层输入减半来突出Self-attention中的主导因子，有效地处理过长的输入序列。
对长时间序列进行一次预测而不是一步步方式进行预测，极大提高了长序列预测的推理速度。

Informer模型的整体框架如下图所示，可以看出该模型仍然保存了Encoder-Decoder的架构：

编码过程（左）：编码器接收长序列输入（绿色部分），通过ProbSparse自注意力模块和自注意力蒸馏模块，得到特征表示。ProbSparse自注意力机制利用稀疏矩阵来替代原来的注意力矩阵，大幅减少算力需求的同时并保持的良好的性能。

解码过程（右）：解码器接收长序列输入（预测目标部分设置为0），通过多头注意力与编码特征进行交互，最后直接预测输出目标部分（橙黄色部分）。这里作者采用的是一次生成式预测方式，并说明该方式相比step-by-step方式推理速度更快且效果相当。

至于特征输入，在 LTSF 问题中，时序建模不仅需要局部时序信息还需要层次时序信息，如星期、月和年等，以及突发事件或某些节假日等。经典自注意力机制很难直接适配，因此Informer提出了三层特征输入形式，如下图所示：

Informer的输入分为了三种位置嵌入表示：

局部时间戳，即Transformer中的固定位置嵌入。
全局时间戳。对于层次时间信息，构建一个词汇表，通过Embedding特征表示每一个词汇。
对齐维度，使用一维卷积将输入的序列标量转化为向量。

模型实践

看完上述的介绍，是否有想一试算法的冲动了呢？笔者也是带着好奇，在幻方萤火集群上尝试复现论文中的实验效果。

作者的开源比较完备，包含ETT（变压器温度）、ECL（耗电量）和WTH（气象）3个数据集，采用PyTorch实现且没有特殊包依赖的模型代码。作者提供了执行的脚本 scripts/*.sh，包括了不同的实验参数：

### M

python -u main_informer.py --model informer --data ETTh1 --features M --seq_len 48 --label_len 48 --pred_len 24 --e_layers 2 --d_layers 1 --attn prob --des 'Exp' --itr 5 --factor 3

python -u main_informer.py --model informer --data ETTh1 --features M --seq_len 96 --label_len 48 --pred_len 48 --e_layers 2 --d_layers 1 --attn prob --des 'Exp' --itr 5

python -u main_informer.py --model informer --data ETTh1 --features M --seq_len 168 --label_len 168 --pred_len 168 --e_layers 2 --d_layers 1 --attn prob --des 'Exp' --itr 5

### S

python -u main_informer.py --model informer --data ETTh1 --features S --seq_len 720 --label_len 168 --pred_len 24 --e_layers 2 --d_layers 1 --attn prob --des 'Exp' --itr 5

python -u main_informer.py --model informer --data ETTh1 --features S --seq_len 720 --label_len 168 --pred_len 48 --e_layers 2 --d_layers 1 --attn prob --des 'Exp' --itr 5

python -u main_informer.py --model informer --data ETTh1 --features S --seq_len 720 --label_len 336 --pred_len 168 --e_layers 2 --d_layers 1 --attn prob --des 'Exp' --itr 5

一共有60组不同的实验，我们用幻方萤火平台快速运行一下。

萤火提供了统一的训练管理平台，可以将大量的训练任务依据优先级负载均衡到不同的GPU上执行，充分利用起计算算力，节省模型训练的时间。我们只需加入几行代码，便可以把训练任务提交给集群，喝杯咖啡的功夫就能拿到结果啦。

登录萤火平台，引入haienv, hfai

# ./main_informer.py
import haienv
haienv.set_env('202111')

# ./exp/exp_informer.py
import hfai

对于每一轮训练，加入接收集群调度的逻辑代码，并做好模型checkpoint的保存

# 获取当前训练的进度
for epoch in range(hfai.get_whole_life_state()//len(train_loader)%self.args.train_epochs, self.args.train_epochs):
    ...
    
    # 对于之前已经训练过的轮次，直接跳过
    if epoch*len(train_loader)+i <hfai.get_whole_life_state():
        continue
    ...
    
    # 状态更新
    steps += 1
    
    # 收到调度信号，保存模型，设置当前执行的状态信息
    if hfai.receive_suspend_command():
        ...
        hfai.set_whole_life_state(steps)
        hfai.go_suspend()

提交任务，等待执行结果

完成

Informer开源的测试数据不大（10MB以内），集群测试中单个Epoch执行花费6-7s左右。对每组实验配一块A100卡，我们很快就可以完成60组实验。在长时序单变量预测和长时序多变量预测这两个任务上，测试结果基本与论文所公布的MSE和MAE结果吻合，在某些参数上甚至跑出了更好的成绩。

体验总结

Informer作为今年AAAI的Best Paper之一，对Transformer模型进行了很多切实有效的改进，使其计算、内存和体系结构更加高效。同时，作者也做了完整的开源，代码结构清晰，笔者能够很流畅的对其进行复现。借助萤火平台，我们很快就拿到了全部的实验数据。综合体验打分如下：

研究指数：★★

该模型着眼于长时间序列的预测问题，是一个广泛研究的基础课题。
开源指数：★★★★★

数据和代码都已开源，代码逻辑清晰可读性高。
门槛指数：★★★★

数据量小，模型计算、内存和体系结构均已优化，普通高性能PC即可运行。
通用指数：★★★

模型是对Transformer的优化，深入改进了自注意力机制，能适用于长序列预测场景，其他场景待验证。
适配指数：★★★★★

依赖简单，PyTorch框架构建的模型，只需要修改几行代码就能在萤火集群上执行。

幻方 AI 紧跟 AI 研究的前沿浪潮，致力于用领先算力助力AI落地与价值创造，欢迎各方数据研究者与开发者们一同共建。

3FS优化 03 | 数据读取模式适配

Qiao — Tue, 15 Sep 2020 13:00:00 GMT

如文章《幻方力量 | 高速文件系列 3FS》中所介绍的，幻方AI设计了一套非常适合深度学习训练的样本读取文件系统，3FS，其采用Direct IO和RDMA Read的读取方式，让模型训练在样本读取部分只用极小的CPU和内存开销，就可以获得超高的读取带宽，从而无需再训练过程中等待加载数据，更充分地利用起GPU的计算性能。

然而，实际应用中会有很多我们之前没有预料到的问题，比如任务间互相影响的问题，主要由任务不同的数据读取模式引起。本文将详细描述该问题的来龙去脉，以及我们的解决思路。

我们知道，幻方AI的深度学习平台提供的是分时调度的深度学习训练服务，用户可以自由地在平台上配置资源，提交任务。对于集群来说，同一时间会有多个任务在计算集群中训练，从存储集群中拉取数据。不同任务的读取模式，处理方式彼此之间会造成一些影响。

我们发现：

如果系统里只有一种任务，即使使用的存储节点再多，每个节点都能读到不错的性能；
如果有不同读取模式的其他任务调度上去跑，即使3FS的总带宽并没有跑满，之前那些任务的读取带宽也会剧烈下降。

如上图所示，计算集群侧有两个计算任务同时从3FS中读取批次数据，训练模型。计算任务1的读取模式只能从3FS存储集群中一部分节点上读取数据，因为集群是负载均衡的，分散到每个节点上读数据大概占其带宽的2/3。计算任务2的读取模式可以从3FS存储集群中大部分节点上读取数据，负载均衡到每个节点上，那些支持了计算任务1的节点只剩下1/3的带宽给计算任务2用，导致那些没有支持计算任务1的节点就算带宽有富余也无法被利用起来。

我们分析了受影响的计算任务1的读取模式，发现它读到的每个样本都很长，大约是 24MB 左右，而且每个 batch 都只从少量文件中读取数据。我们默认的块大小（chunk size）是 2MB，这个任务一个 batch 大约只从40个左右的存储服务节点中同时读取数据，而其他的类似计算任务2的任务几乎都能从接近60个，甚至是所有64个存储服务节点中同时读取数据。在单台存储服务节点总带宽有限的情况下，一个任务连接的存储服务节点少，则分到每单台存储节点的带宽需求就多。当多个任务在请求3FS读取数据时，满负荷下会造成部分任务的需求带宽无法满足的情况，进而影响了计算节点侧的数据读取性能。

因此，对于任何任务从3FS上读取数据，要尽可能让64个存储服务节点能同时读取数据，进而降低单存储节点的带宽需求。我们把块大小缩小到 512KB，并且增加了一个文件分布到的服务器数目的控制，最终成功解决了这个问题。

FFRecord | 高性能训练数据格式

小杞 — Mon, 27 Jul 2020 21:00:00 GMT

在之前的文章《幻方萤火 | 高速文件系统 3FS》中提到，幻方AI自研了一套适合深度学习模型训练场景的文件读写系统，3FS，能提供高性能的批次数据读取，提高模型的训练效率。对于用户而言，使用3FS非常简单，只需要掌握我们封装设计的高性能数据格式，FFRecord，将数据存入萤火集群中即可。

那么FFRecord和一般的数据格式有什么不同？它应该要有哪些设计要求呢？本期文章将为大家分享FFRecord设计背后的故事，讲述幻方AI设计FFrecord的理念，展现FFRecord的高性能和便捷。

项目地址：https://github.com/HFAiLab/ffrecord

转换指引：https://github.com/HFAiLab/ffrecord_converters

概述

我们知道，3FS专门针对模型批量读取样本数据这个场景进行了深度的优化，和一般的文件系统不同，3FS 文件系统有如下的一些特点：

大量打开、关闭小文件的开销比较大；
支持高吞吐的随机批量读取。

如果要充分利用 3FS 文件系统的高效读取性能，我们希望读取的样本数据格式应该满足以下两个条件：

能够知道每一条样本的文件偏移量，方便做随机批量读取；
避免大量的小文件，所有的样本数据在一个或多个大文件里。

然而，目前网上的开源文件格式，比如 TFRecord 等并不完全满足上面的两个条件。TFRecord 能够避免大量的小文件，但是只支持顺序读，而且每次只能读取一条样本。它在不需要随机访问样本的情况下能够获得优秀的性能；但在很多机器学习的场景下，用户训练模型时每一步需要随机采样多个样本。

为此，我们专门设计了 FFRecord (Fire Flyer Record) 文件格式，其能够充分利用 3FS 文件系统的高效读取性能，其包括如下优势：

合并多个小文件，减少了训练时打开大量小文件的开销，对存储后端更加友好；
支持随机批量读取，提升读取速度；
包含数据校验，保证读取的数据完整可靠。

以下是 FFRecord 文件格式的存储 layout:

+-----------------------------------+---------------------------------------+
|         checksum                  |             N                         |
+-----------------------------------+---------------------------------------+
|         checksums                 |           offsets                     |
+---------------------+---------------------+--------+----------------------+
|      sample 1       |      sample 2       | ....   |      sample N        |
+---------------------+---------------------+--------+----------------------+

在 FFRecord 文件格式中，我们会把每一条样本的数据做序列化按顺序写入，同时我们会存储每一条样本在文件中的偏移量和对应的 crc32 校验和，方便我们做随机读取和校验数据。相比于 TFRecord 等格式，我们没有直接给用户提供序列化和反序列化的接口，而是把这一步交给用户，用户可以用自己喜欢的方式 (比如 pickle)去做序列化或者反序列化。

FFRecord 转换规则

转换文件数量和大小

为满足最佳性能，FFRecord 的转换需要满足如下两个要求：

每个文件不小于 256MB （512 x 512KB）
满足条件1的情况下，文件数量尽可能大于100，小于200

转换方式对比

FFRecord 中的数据通过二进制格式进行存储，有如下两种将图片数据转换为二进制数据的方法：

方法 1：通过 pickle 转换为二进制

image = Image.open(fname)
data = pickle.dumps(image)

方法 2：使用二进制格式读取文件

with open(img_file, "rb") as fp:
    img_bytes = fp.read()

两种方法的性能性能对比：

速度对比：方法 1 (pickle) 比方法 2 (二进制文件) 读取速度更快
- 方法 1 读取速度大约为 10 GB/s，方法 2 读取速度大约为 1GB/s
存储大小对比：方法 2 (二进制文件) 转换后文件体积更小
- 方法 1 的存储体积约为方法 2 的 5 倍

数据封装

接下来我们介绍一下 FFRecord 的封装方式。FFRecord 文件的读写分别可以通过 FileReader 和 FileWriter 对象来完成。以下是一个简单的使用示例：

数据写入

from ffrecord import FileWriter

samples = [i for i in range(100)]  # 准备写入的样本
fname = 'test.ffr'
n = len(samples)                   # 样本的数量
writer = FileWriter(fname, n)

for i in range(n):
    data = pickle.dumps(samples[i])  # 序列化
    writer.write_one(data)

writer.close()

FileWriter.write_one 方法会把传入的数据写到文件的末尾，然后在文件的头部写入这条数据在文件中的偏移量和对应的 crc32 校验和

数据读取

from ffrecord import FileReader

fname = 'test.ffr'
reader = FileReader(fname, check_data=True)
print(f'Number of samples: {reader.n}')

indices = [3, 6, 0, 10]      # 样本的索引
data = reader.read(indices)  # 返回样本的原始数据
samples = [pickle.loads(x) for x in data]  # 反序列化

FileReader.read 方法会先根据传入的索引找到每条样本在文件中的偏移量并计算出占用的字节数，然后调用 Linux Asynchronize I/O (AIO) 的接口同时处理多个样本的读取请求，把结果返回给用户。

Dataset 和 DataLoader

PyTorch 本身提供了 map style 的 Dataset 和 DataLoader 的接口，然而它们的接口对需要批量读取的场景不是很友好，PyTorch Dataset 需要实现 __getitem__ 方法，该方法传入一个样本的索引并返回一条样本，这样如果我们希望批量读取多个样本就会比较困难。为了让 PyTorch 用户更加方便的使用 FFRecord 读取训练样本，我们提供了定制的 FFDataset 和 FFDataLoader。FFDataset 同样需要实现一个 __getitem__ 方法，但这个方法接受一个 batch 的样本索引作为输入，返回一整个 batch 的样本。FFDataset 可以搭配 FFDataLoader 一起使用，实现多进程异步加载数据。

以下是一个简单的使用示例：

from ffrecord.torch import Dataset as FFDataset
from ffrecord.torch import DataLoader as FFDataLoader

class ImageNet(FFDataset):
    def __init__(self, fname, transform=None):
        self.reader = FileReader(fname, check_data=True)
        self.transform = transform

    def __len__(self):
        return self.reader.n

    def __getitem__(self, indices, data):
        data = self.reader.read(indices)
        samples = []
        for bytes_ in data:
            img, label = pickle.loads(bytes_)
            if self.transform:
                img = self.transform(img)
            samples.append((img, label))
        return samples

dataset = ImageNet(fname, transform)
loader = FFDataLoader(dataset, batch_size=32, shuffle=True, num_workers=8)

性能测试

我们在 3FS 存储集群上进行了性能测试，以读取 ImageNet 数据集为例，比较如下两种数据读取方式的性能：

大量小文件 + PyTorch DataLoader；
FFRecord + FFDataLoader。

为了模拟在分布式训练模型时候的读取模式，我们会起 8 个互相独立的进程，每个进程有自己的 DataLoader，然后每个 DataLoader 会再开 16 个子进程 (num_workers = 16)，batch size 设为128，然后我们统计每个 DataLoader 读取 1000 个 batch 的时间并取平均。

上图展现了我们实验得到的结果，从图中我们可以看到使用了 FFRecord + FFDataLoder 后速度能够得到明显的提升。

hfreduce | 高性能的多卡并行通信工具

High-Flyer — Wed, 04 Mar 2020 21:00:00 GMT

hfreduce 是幻方AI自研的高性能的多卡并行通信工具，是幻方“萤火二号”计算存储分离后，计算服务中的重要一环，根据萤火二号的集群特性专为计算节点设计的 allreduce 工具。本质上 hfreduce 相当于 PyTorch 中的 DistributedDataParallel（DDP），只不过使用 CPU 做加法运算以计算总梯度，而不是调用其他的集体通信库（CCL），比如说 NCCL，传递梯度到不同的显卡上，再各自计算总梯度。

为什么要这么设计呢？本期文章将为大家解密，详细介绍幻方AI设计 hfreduce 的背景与理念。文章会阐明 hfreduce 对模型训练整体加速的效果以及具体适用的场景。

背景

幻方AI之所以采用 CPU 来做加法运算，计算总梯度，是因为萤火二号上的计算节点配置比较特殊。

幻方的主要AI场景是金融行为分析、自然语言处理、生物分子结构预测等。在这些场景中，基本是数据规模大而模型大小适中。换句话说，在A100显卡40G的显存中，完全可以装得下一个完整的模型和批次样本数据。因此，模型的加速主要是依赖大量的数据并行，让尽可能多的显卡参与训练，再同步梯度。

正是因为上述数据并行的应用场景，随着集群规模不断变大，我们发现使用 NCCL 进行多 GPU 之间的集体通信，梯度信息多次在 PCIe 上遍历就会成为一个性能瓶颈。此外，我们无法使用在交换机上做加法运算计算梯度的方法，因为萤火二号集群中会同时运行大量的训练任务，它们各自需要做reduce以计算总梯度，而交换机之前并不支持在同一个网络内进行如此多路的reduce计算。

框架设计

近些年深度学习模型逐渐尝试更大的参数空间来处理大规模的数据，使用多个GPU一同加速训练成为很多AI从业者的必然选择。NCCL是英伟达发布的用于实现多GPU之间集体通信的工具，在PCIe，NVLink上可以实现较高的通信速度。NCCL借助CPU，可以传播，计算不同显卡上的梯度信息，对并行训练中的多个显卡进行梯度信息的同步。在有NVLink的条件下，在同节点不同显卡内做reduce计算对 NCCL 来说轻而易举，但如果没有NVLink，梯度信息多次在 PCIe 上遍历就会成为一个性能瓶颈。

既然我们是因为没有 NVLink 而放弃了 NCCL，那么把梯度数据都先拷贝到主内存里然后在 CPU 上做加法运算计算总梯度也就是很自然的事了，因为这样可以避免梯度数据在多个 GPU 之间进行多次拷贝。hfreduce 首先从所有并行的GPU里一次性收集梯度信息，然后在CPU里进行reduce计算，完成后将计算出来的梯度一次性广播到所有 GPU，从而将 PCIe 流量减半。此外，CPU 只用于将所有梯度信息相加汇总，与 NCCL 不同，免去了在 GPU 上再计算梯度的开销。

如上图所示，与 NCCL 一样，我们也采用两棵树的拓扑结构进行节点间的通信。对于节点比较多的情况，它能比环形拓扑更好，因为节点之间的对数距离会更短。我们设计让树结构更加平衡一些，所以理论上它能比 NCCL 的树结构稍微快一些。

在CPU 上的加法运算也进行了精心设计。所有加载到内存中的梯度信息都在 NUMA 节点之间进行交互，绑定到 CPU 内核的线程从同一个 NUMA 节点中挑选出梯度数据并将它们相加。 AUX 256 指令优于 AUX 512 指令，因为 AMD CPU 不支持后者，如果所有内核都在执行指令，则会导致 CPU 频率受限。

由于所有的加法运算都是在 CPU 内核上执行的，我们观察到内存带宽是瓶颈，而不是CPU的算力。因此，建议使用更大的样本批次，因为读取它们可能会减少内存带宽的竞争。 PCIe 的宽松排序可能会提高Reduce的性能，但实际上更快的Reduce可能会减慢样本读取的速度，从而减慢整个训练的速度。 Reduce的性能很重要，但就目前来说，该方案仍然只是一个局部最优解。

使用案例

我们给hfreduce设计了Python的接口，使其可以轻松地加入 PyTorch 的训练中。该库能自动处理节点地址交换、梯度归零和梯度的最终平均。更重要的是，该库能在模型代码中直接嵌入，以便在计算梯度后立即开始Reduce计算。 hfreduce的使用基本上就像创建一个 Reducer 实例一样简单。

具体的，在萤火二号上训练模型时，可以参考如下案例代码：

import hfreduce.torch as hfr

model = Model().cuda()
loss = Loss(model)
optimizer = Optimizer(model)

reducer = hfr.AsyncReduceFloat(registry_host, port, proc_rank, procs_per_node, node_rank, node_cnt, model)

for sample in sampels:
    reducer.zero_grad() # instead of optimizer.zero_grad(), call reducer.zero_grad()
    
    model(sample)
    loss.backward()
    
    reducer.synchronize()
    
    optimizer.step()

性能测试

我们使用 PyTorch 内嵌的 VGG16 模型进行实验，测试hfreduce的性能。

如上图所示，hfreduce整体耗时更少，明显比没有NVLink的NCCL更快。使用 PyTorch 进行模型训练，即使在后台使用 PyTorch DDP 调用 NCCL 进行 allreduce，hfreduce 仍然胜过 NCCL。

3FS优化 02 | 客户端内存使用优化

Qiao — Wed, 04 Dec 2019 12:00:00 GMT

如文章《幻方力量 | 高速文件系列 3FS》中所介绍的，幻方AI设计了一套非常适合深度学习训练的样本读取文件系统，3FS，其采用 Direct IO 和 RDMA Read 的读取方式，让模型训练在样本读取部分只用极小的CPU和内存开销，就可以获得超高的读取带宽，从而无需再训练过程中等待加载数据，更充分地利用GPU的计算性能。

我们知道，文件系统一般分为客户端与服务端。在3FS文件系统中，客户端部署在计算节点侧，通过网络访问部署在存储节点侧的3FS服务端。这个过程中有很多问题值得优化。本期文章将和大家聊聊，客户端的内存使用优化的问题。

幻方AI的深度学习平台是计算存储分离的，计算集群只负责计算，无论是数据读取、模型加载还是梯度 reduce 等都需要通过网卡传入内存再传入显存，中间需要进行多次内存拷贝，对内存带宽的使用非常严重。如果能有效优化内存的使用，将对模型训练的性能带来极大的提升。

概述

3FS的客户端是指部署在计算节点侧的3FS，其作用是与服务端的3FS通过网络联通，从服务端请求数据并解析。对于幻方萤火二号集群来说，集群里的每台机器上有一张高速网卡，对于一般的文件系统，我们来看一下每秒钟读取23.5GB数据所需的内存开销：

从网卡读到内核态内存：一次内存写入
从内核态读到 Page Cache，一次内存读取一次内存写入
从 Page Cache 读到用户态内存：一次内存读取和一次内存写入
另外，由于内核里面不支持Non-temporal store的寄存器指令，写入时还需要多出一到两倍的内存读取。详情参见这里。

如上所述，如果我们想要达到23.5 GBps的读取性能，需要消耗每秒 23.5 * 6~7 = 141~165GBps 的内存带宽。

而现在市面上最先进的训练平台，比如双路AMD Epyc Rome/Milan的总内存带宽在270GBps~330GBps之间，双路Intel Cascade Lake平台的总内存带宽更是只有220GBps左右，也就是说，对于一般的文件系统，光是读取数据所消耗的内存带宽就要占掉机器总内存带宽的50%~70%！而我们在使用这些机器的时候，显然只有数据是不够的，如果IO占了这么多带宽，那么留给计算和通信的资源就少了。

因此减少内存拷贝操作、降低内存开销，是我们优化3FS客户端性能的重点方向。我们采用RDMA Read的技术来解决该类问题。

RDMA Read

我们通过 RDMA Read 到用户态内存的方式来进一步避免内核态到用户态的内存拷贝。如前文给出的例子，使用 RDMA Read 之后，读数据的操作变为一步：从网卡直接读到用户态内存，也就是说，读满 23.5GBps 的网卡，只需要消耗 23.5GBps 的内存带宽，占总内存带宽的 10% 不到。

我们选择了计算节点侧发起 RDMA Read ，而不是从存储节点侧做 RDMA Write。这是因为3FS的批量读取导致同一网络上的读取请求可能对应了很多应用的 Buffer，而 RDMA Write 不支持分发到接收方的多个地址，这意味着在单个读取长度很小的时候，会产生很多很小的 RDMA 请求，以至于显著影响读取的带宽。

然而需要注意的是，读取的数据要做对齐，RDMA Read 并不是简单地把读到的数据填到用户的 Buffer 里，而是需要把存储节点侧对齐用的数据读到内核中分配的 padding buffer 里，在计算节点侧对齐后再丢弃掉。在测试性能的过程中我们发现，不同大小的 RDMA Read 请求会互相影响：

每个请求过小，比如只读几个 KB，效率不高；
每个请求太大，一次读取几个 MB，在网络繁忙的时候效果也并不算好。

我们把小的请求合并起来，把大的请求拆开，最后定在每个 RDMA 读64KB 左右，收获了不错的效果。

小结

本文介绍了幻方AI在设计3FS文件系统时，针对客户端内存使用优化的一些思考。我们采用 Direct IO 和 RDMA Read 的读取方式，让服务端的数据通过网卡直接加载到用户态内存中，减少内存带宽的占用，让模型训练在样本读取部分只用极小的CPU和内存开销，就可以获得超高的读取带宽，从而无需再训练过程中等待加载数据，更充分地利用GPU的计算性能。

3FS能跑出这样的性能也不是一蹴而就的，幻方AI在长期实践中对其不断优化，遇到了不少挑战。本系列文章将陆续为大家分享一些有关3FS性能优化的小故事，作为抛砖引玉，其中还不乏我们踩过的坑。欢迎业界的学者专家们一起探讨。

3FS优化 01 | 服务端优化

Qiao — Sun, 17 Nov 2019 11:00:00 GMT

如文章《幻方力量 | 高速文件系列 3FS》中所介绍的，幻方AI设计了一套非常适合深度学习训练的样本读取文件系统，3FS，其采用 Direct IO 和 RDMA Read 的读取方式，让模型训练在样本读取部分只用极小的 CPU 和内存开销，就可以获得超高的读取带宽，从而无需再训练过程中等待加载数据，更充分地利用起 GPU 的计算性能。

我们知道，文件系统一般分为客户端与服务端。在 3FS 文件系统中，客户端部署在计算节点侧，通过网络访问部署在存储节点侧的 3FS 服务端。这个过程中有很多值得优化的问题。本期文章将和大家聊聊，服务端的优化问题。

为了减少内存拷贝，提高读取吞吐，3FS 服务端的主要操作是：

使用 Direct IO 来做读取，减少了磁盘到 Page Cache 的拷贝
使用 AIO 接口提高读取吞吐
在客户端进行数据对齐，简化接口使用

Direct IO

Linux 提供两种文件 IO 模式，分别是 Buffered IO 和 Direct IO。Buffered IO 使用内核提供的 Page Cache 缓存磁盘上的数据，可以为多次访问的数据带来显着的性能优势。但对深度学习训练来说，每一次迭代读取的都是新数据，缓存失去意义，Page Cache 反而成为了负担。所以 3FS 服务端会默认使用 Direct IO 进行数据读取，绕开 Page Cache 以减少一次数据拷贝，直接从磁盘读数据到用户态内存，同时也规避了 Page Cache 的高内存占用。

AIO

AIO 是 Linux 提供的异步 IO 接口，它支持应用提交多个异步 IO 操作在后台并发执行，并在前端同步地等待操作完成。AIO 目前仅支持以 Direct IO 方式打开的文件。相较于传统的同步 IO 接口，AIO 可以在更低的 CPU 占用下提供更高的读写吞吐。幻方 AI 的深度学习训练框架中默认使用 AIO 接口进行训练样本数据的读取。此外，对于读取大小超过阈值的同步读取操作（阈值默认 1MB），3FS 的客户端会将其转为 AIO 操作以提高读取性能。同时我们也正在适配内核新提供的 io_uring 接口，相信可以带来进一步的性能提升。

对齐

Direct IO 在使用上有一个很大的限制，它需要 buffer 的内存地址、每次读取的 size 以及 offset 均与底层设备的逻辑块大小对齐，一般是 512 字节对齐。为了提高接口的易用性，3FS 的客户端会对 AIO 读取操作进行必要的对齐，它会在 buffer 的头部和尾部增加额外缓存以实现 512 字节对齐，并在 RDMA Read 接收数据时忽略掉头部和尾部。这样就实现了以下两种效果：

用户的 AIO 操作不再需要对齐，使用上更简单
用户的同步 IO 操作可以在客户端直接转为 AIO 操作而不需要考虑对齐和 O_Direct，提高性能

小结

本文介绍了幻方AI在设计3FS文件系统时，针对服务端优化的一些思考。我们采用 Direct IO 和异步对齐的读取方式，让服务端的数据加载更符合模型训练的使用场景，从而获得更好的读取性能。

幻方力量 | 高速文件系统 3FS

High-Flyer — Thu, 13 Jun 2019 21:00:00 GMT

3FS是幻方AI自研的高速文件系统，是幻方“萤火二号”计算存储分离后，存储服务中的重要一环，全称是萤火文件系统（Fire-Flyer File System），因为有三个连续的 F，念起来不是很容易，因此被简称为 3FS。

3FS 是一个比较特殊的文件系统，因为它几乎只用在AI训练时计算节点中的模型批量读取样本数据这个场景上，通过高速的计算存储交互加快模型训练。这是一个大规模的随机读取任务，而且读上来的数据不会在短时间内再次被用到，因此我们无法使用“读取缓存”这一最重要的工具来优化文件读取，即使是超前读取也是毫无用武之地。因此，3FS的实现也和其他文件系统有着比较大的区别。

本期文章将为大家解密幻方AI是如何设计与实现3FS的，以及最终能实现的模型训练加速效果。

硬件设计

3FS文件系统整体的硬件设计如下图展示：

可以看到，3FS文件系统包含两大部分：数据存储服务与高速交换机。数据存储服务与计算节点分离，专门用于存储模型训练需要用到的样本数据。每个存储服务节点有16张各15TB的SSD硬盘和2张高速网卡，读取性能强劲，网络宽带强大。

3FS节点与计算节点（Clients）通过800口高速交换机进行连接。值得注意的是，因为一台交换机要连接 600 台左右的计算节点，每个计算节点就只能使用一张网卡，这样，从3FS上读取样本数据的流量和训练产生的其他数据流量（梯度信息、数据并行信息等）就不得不分享这一张网卡的带宽，这对3FS整体读取性能带来了一些挑战。

软件实现

前文中提到，3FS解决的是模型训练中样本数据读取这个特定的使用场景。与常见的文件读取方式不同，训练需要读取的样本是随机的，而且往往一个batch里用到的是互不相关的样本。注意到这一点后，我们决定设计异步的文件读取方式。

具体的，如上图所示，3FS使用基于Linux的 AIO 和 io_uring 接口来完成样本的读取，因为在3FS的场景下，File Cache 完全没有作用，反而会以一种用户很难控制的方式消耗系统的内存，影响后续任务的运行，因此我们关闭了 File Cache，只用 Direct I/O 模式做数据的读取。但要注意的是，用这种方式读取，buffer 指针，offset 和长度都是需要对齐的，如果让用户来做这个对齐，就会产生额外的内存拷贝，因此我们在文件系统内部把对齐做掉了，既优化了性能又方便了用户。

用户使用3FS非常简单，只需要将样本数据转化成FFRecord格式，存入3FS中即可。FFRecord格式是幻方AI开发的一种适用于3FS性能的二进制序列存储格式，适配了 PyTorch 的 Dataset 和 Dataloader 接口，可以非常方便的加载发起训练。具体项目地址： https://github.com/HFAiLab/ffrecord

在使用幻方萤火进行模型训练时，只需要将原始数据做完特征工程，转化成输入模型的样本数据，通过3FS加载，你就可以享受到极致的存储性能体验。

压力测试

幻方萤火二号目前有64台存储服务器构成3FS文件系统。想象一下，训练一个 ResNet，读取 ImageNet 数据，ImageNet 压缩文件大小148G，解压整理成 FFRecord 格式的二进制训练样本数据有700G+，假设 batch_size 为400读满单个A100 40G显存，使用3FS读取数据在最优性能下每 Epoch 读取 ImageNet 数据只需耗时0.29s~0.10s左右。这极大降低数据加载的开销，充分利用GPU的计算时间，提高GPU利用率。

上图展示了分布式训练 ResNet 实际每轮 Epoch 耗时情况。可以看到，即使在集群满负荷状态下，读取数据的耗时仅占总耗时的1.8%左右，具有相当不错的数据读取性能。