超越GPT-4!腾讯AI新研究打破长文本生成模型限制,序列并行技术再突破
引言:探索长文本生成模型的新趋势
在人工智能领域,随着生成模型中上下文长度的不断增长,我们见证了一种趋势的形成。从Claude在大型语言模型(LLMs)中将序列长度扩展到100K标记,到OpenAI的GPT-4将上下文长度扩展到128K标记,再到多模态模型的出现,如Gemini 1.5 Pro拥有高达1000万标记的上下文长度,以及OpenAI的Sora模型支持至少100万视觉标记。这些突破性进展强调了生成AI技术需要巧妙处理更大上下文长度的重要性。
序列并行(Sequence Parallelism,简称SP)作为一种划分输入序列的技术,已经成为训练或推理更长序列的有希望的方法。通过两年的初步探索,到2023年下半年,两项标志性的工作——DeepSpeed-Ulysses和Ring-Attention——标志着SP技术的成熟。DeepSpeed-Ulysses在序列长度和计算设备成比例增加时保持恒定的通信量,而Ring-Attention通过计算和通信的重叠隐藏了SP引入的P2P通信成本。然而,挑战仍然存在,例如DeepSpeed-Ulysses的SP并行度受到注意力头数的限制,Ring-Attention由于矩阵乘法的细分而降低了计算效率。这些限制目前阻碍了序列并行在分布式Transformer计算中的更广泛应用。
论文标题、机构、论文链接和项目地址
论文标题: A Unified Sequence Parallelism Approach for Long-Context Generative AI
机构: Tencent
论文链接: https://arxiv.org/pdf/2405.07719.pdf
项目地址: https://github.com/feifeibear/long-context-attention
序列并行技术(SP)的概述
序列并行技术(Sequence Parallelism,简称SP)是一种将输入张量的序列维度分割到多个计算设备上的技术,这种技术对于解锁生成性AI模型的长上下文能力至关重要。SP通过在多个计算设备上分配输入序列,使得每个设备处理输入序列的一个子段,从而支持更长的序列处理能力。
近年来,随着生成性AI模型上下文长度的不断增长,SP技术的重要性日益凸显。例如,Claude在大型语言模型(LLMs)中将序列长度扩展到了100K个标记,而OpenAI的GPT-4则将上下文长度扩展到了128K个标记。此外,多模态模型的出现进一步推动了这一趋势,例如Gemini 1.5 Pro拥有高达1000万个标记的上下文长度。
SP技术的关键在于它能够有效地分配和管理大规模输入序列的计算任务,从而使得模型能够处理更长的序列而不会受到单个设备内存限制的束缚。这种技术的发展经历了初期的探索,并在2023年后期通过DeepSpeed-Ulysses和Ring-Attention两个里程碑式的工作达到成熟。
深入解析 DeepSpeed-Ulysses 和 Ring-Attention
1. DeepSpeed-Ulysses(SP-Ulysses)
DeepSpeed-Ulysses是一种序列并行方法,它通过All2All通信操作处理分割后的Q(查询)、K(键)、V(值)和O(输出)张量。这种方法的特点是当序列长度和计算设备成比例增加时,通信量保持不变。在All2All操作后,这四个张量的分割从序列维度L转移到注意力头数维度hc。这样,每个注意力头的softmax(QK^T)V计算可以完整地进行,而不会因为张量的分割而中断。
2. Ring-Attention(SP-Ring)
Ring-Attention是另一种序列并行方法,它可以看作是FlashAttention的分布式版本。在计算输出张量O的各个块时,如果所需的K和V块不在本地可用,则使用点对点(P2P)通信从其他设备获取。这种方法的通信可以以环形方式组织,每个设备同时发送和接收K、V块,允许通信与计算重叠。
这两种方法各有优势和局限。DeepSpeed-Ulysses对注意力头数敏感,其并行度受到注意力头数的限制,而Ring-Attention在计算和通信效率上存在挑战。尽管如此,这两种技术提供了处理长序列的有效途径,为未来的研究和应用奠定了基础。
提出统一的序列并行方法
在探索序列并行技术(SP)的发展过程中,DeepSpeed-Ulysses和Ring-Attention作为两种主要的实现方式各有优势和局限。DeepSpeed-Ulysses在增加计算设备和序列长度时能保持恒定的通信量,而Ring-Attention通过计算和通信的重叠隐藏了由SP引入的点对点(P2P)通信成本。然而,DeepSpeed-Ulysses的并行度受到注意力头数的限制,而Ring-Attention在分块矩阵乘法中的计算效率较低。
为了克服这些限制,我们提出了一种统一的序列并行方法,将DeepSpeed-Ulysses和Ring-Attention的技术结合起来。这种统一方法通过在序列维度上的混合并行策略,允许在不同的计算设备上分别运行DeepSpeed-Ulysses和Ring-Attention,从而优化了通信模式和计算效率。
具体来说,统一的SP方法在一个2D网格的进程组中运行,其中DeepSpeed-Ulysses沿着网格的行进行操作,而Ring-Attention沿着列进行。这种设置不仅提高了模型架构和网络硬件的适应性,还通过更有效的负载平衡和通信模式,提高了整体的计算性能。
与其他并行技术的比较
在并行计算技术中,数据并行(DP)、张量并行(TP)、ZeRO优化和流水线并行(PP)是常见的方法。每种技术都有其特定的优势和应用场景,而序列并行(SP)作为一种新兴的并行方法,其与这些现有技术的关系和综合应用同样重要。
1. 数据并行(DP)与序列并行(SP)数据并行主要通过在多个计算设备上复制模型来并行处理不同的数据批次。SP与DP在激活的通信成本上相当,但SP在注意力模块中引入了额外的通信开销。当使用Ring方法时,尽管额外的P2P通信可以与计算重叠,但仍可能引入性能问题。
2. 张量并行(TP)与序列并行(SP)张量并行通过在多个计算设备上分片模型参数来实现并行。与SP相比,TP在通信成本上具有优势,尤其是在处理长输入序列时。然而,使用MQA技术可以显著降低SP的通信成本,使得在某些情况下SP可能比TP更有优势。
3. ZeRO优化与序列并行(SP)ZeRO通过在多个设备上分片优化器状态、梯度和参数来减少每个设备的存储需求。SP可以与ZeRO一起在同一个进程组中运行,这使得SP在内存成本上可以与TP相媲美,特别是当使用ZeRO-3时。
4. 流水线并行(PP)与序列并行(SP)流水线并行通过在不同的计算设备上分配模型的不同层来实现并行,而SP则是在变压器块内部分割张量。因此,SP与PP是完全兼容的,可以在4D混合并行系统中共同使用。
通过这些比较,我们可以看到统一的序列并行方法不仅提高了模型的可扩展性和效率,还能与其他并行技术有效结合,为处理大规模和长上下文的生成模型提供了强大的支持。
实验验证和性能分析
1. 实验设置和结果
在实验中,我们使用了两个8xA800 GPU节点来测试Unified SP方法在不同序列长度下的性能。根据表3和表4的数据,我们发现在使用8xA100-SXM4 NVLink节点时,当ulysses_degree设置为8时,无论序列长度是32K还是128K,都能达到最高的吞吐量。这验证了SP-Ulysses在处理大规模序列时的优势。然而,在8xL20 PCIe节点上,最佳的设置是ulysses_degree为4,ring_degree为2,这表明在不同的硬件配置下,最优的并行度设置可能会有所不同。
2. 性能对比
我们进一步将SP-Unified方法与Megatron-LM中的现有SP实现进行了比较。结果显示,尽管Megatron-DeepSpeed的SP-Ulysses实现在性能上落后于Megatron-LM,但我们的Unified SP方法在集成了负载平衡策略后,显示出了更优的性能。特别是在处理大规模序列任务时,Unified SP的表现优于单独使用SP-Ulysses或SP-Ring的策略。
序列并行技术在大规模集群上的应用前景
1. 大规模应用潜力
随着序列长度的不断增加,传统的并行技术已经无法有效地处理超大规模的序列数据。Unified SP方法通过结合SP-Ulysses和SP-Ring的优点,不仅提高了计算效率,还降低了通信成本,使其成为处理大规模序列的理想选择。此外,该方法对网络硬件的要求更为灵活,适用于各种不同的网络拓扑结构,这对于在大规模集群上部署模型尤为重要。
2. 未来研究方向
尽管当前的实验结果已经很有希望,但在将序列并行技术应用于超过10,000个GPU的大规模训练任务时,仍有许多挑战需要克服。例如,如何进一步优化通信策略以适应不同的网络条件,以及如何处理更长的序列长度以提高模型的表现。此外,与ZeRO-3和MoE等其他并行技术的结合也是未来研究的重要方向,这可能会开辟新的可能性,以支持更复杂和更大规模的AI模型训练。
总结
在本文中,我们提出了一种将DeepSpeed-Ulysses和Ring-Attention结合的序列并行方法。这种方法整合了两种技术的优势,扩展了其适用性,并在某些情况下提供了更优越的性能。我们系统地分析了序列并行与其他已建立的并行方法之间的相互作用,并得出了一系列最佳实践结论。这些结论通过在两个GPU节点上获取的实验数据得到了验证,证实了我们的发现在实际应用中的影响。
通过深入探讨序列并行技术,我们不仅解决了长序列输入在分布式Transformer计算中的挑战,还通过提出统一的序列并行方法,优化了通信和内存成本。此外,我们的研究强调了序列并行与数据并行、张量并行和管道并行等其他并行策略的兼容性和互补性,为设计混合4D并行系统提供了策略和见解。
实验结果显示,统一序列并行方法在处理大规模语言模型时,尤其是在长序列上,不仅提高了计算效率,还优化了网络通信和内存使用。这一成果为未来在更大规模的集群上部署高效的大模型训练提供了可能。
我们对序列并行技术的未来发展持续保持乐观态度,并期待其在推动生成式AI模型处理更长上下文长度方面发挥更大作用。同时,我们也认识到,随着模型和数据规模的不断扩大,未来的研究需要进一步探索序列并行技术与其他先进并行技术如专家模型并行的结合,以实现更高效的训练过程和更优越的模型性能。
本文转载自AI论文解读,作者: 柏企