YunChang：一种统一的序列并行（USP）注意力机制，用于长上下文大模型的训练和推理。

[技术报告] USP：一种适用于长上下文生成式 AI 的统一序列并行方法

本仓库提供了一种序列并行方法，它融合了两种流行的分布式注意力机制——DeepSpeed-Ulysses-Attention 和 Ring-Attention 的优势，从而具备更通用、更强的适应性以及更好的性能。该项目基于 zhuzilin/ring-flash-attention 构建，并参考了 DeepSpeed-Ulysses 的实现。

USP 已被应用于 NVIDIA/TransformerEngine 的 AttnFuncWithCPAndKVP2P 中。您可以在 API attn_forward_func_with_cp 中使用它。

为什么不能单独使用 Ulysses 和 Ring Attention？

• Ulysses 对注意力头的数量较为敏感。其并行度不能超过头数，因此不适用于 GQA（分组查询注意力）和 MQA（多查询注意力）场景。例如，当只有一个注意力头时，Ulysses 无法有效工作。此外，由于张量并行也需要在头数维度上进行划分，因此要实现 Ulysses 与张量并行的兼容性较为困难。

• 在计算和通信效率方面，Ring-Attention 低于 Ulysses。Ring-Attention 将 Query、Key 和 Value（QKV）分割成较小的块，这可能导致在使用 FlashAttention 时效率下降。即使通信和计算过程完全重叠，总执行时间仍会落后于 Ulysses。此外，Ring-Attention 使用异步点对点通信，不仅带宽利用率低于集体通信方式，而且在大规模部署中还存在潜在的通信死锁风险。

LongContextAttention，又称统一序列并行和混合序列并行

LongContextAttention 是一种 统一的序列并行，也称为 混合序列并行，它结合了 DeepSpeed-Ulysses-Attention 和 Ring-Attention 的优点，从而克服了这两种方法各自的局限性。

1. 安装

FlashAttention 是最重要的外部依赖项，也是安装和使用 yunchang 时经常出错的原因。Yunchang 支持 flash_attn 2.6.x 和 2.7.x 版本，包括 v3 和 v2 版本。此外，yunchang 还支持在没有 flash_attn 的情况下运行，这非常适合 NPU 设备。

如下图所示，根据是否使用 flash_attn，有三种使用方法：

1. 对于 H100、B100 等支持 FA v3 的硬件，ring_flash_attn 使用 FA v3。

2. 对于 A100、L40 等支持 FA v2 的硬件，ring_flash_attn 使用 FA v2。

3. 对于 NPU 等不支持 FA 的硬件，可以使用 PyTorch 实现注意力计算。在这种情况下，无需安装 flash_attn，应使用 LongContextAttention(ring_impl_type="basic", attn_type=AttnType.TORCH_EFFICIENT)。注意：AttnType.TORCH_EFFICIENT 不支持反向传播。

选项 1：通过 pip 安装

pip install flash-attn

pip install yunchang

应用 FlashAttention V3：由于 FA V3 处于测试版发布阶段，您需要从源代码安装 FlashAttention V3。

请按照 FlashAttention 测试版发布 的说明，为 NVIDIA Hopper GPU 安装 V3 版本。我们采用了 2024 年 11 月 10 日的提交 b443207c1fc4c98e4532aad4e88cfee1d590d996。

选项 2：本地构建

pip install .

AMD GPU 的安装：install_amd.md

2. 使用

请参阅 test/test_hybrid_qkvpacked_attn.py 和 test/test_hybrid_attn.py 了解具体用法。

简而言之，我们以 zigzag 式的环形注意力实现为例：

1. 调用 set_seq_parallel_pg 设置进程组。
2. 使用 zigzag_extract_local 提取局部张量。为了使环形注意力负载均衡，我们需要对输入 token 或输入张量进行重新排序。
3. 然后将 LongContextAttention(ring_impl_type="zigzag") 作为注意力实现的直接替代品使用。

from yunchang import (
    AsyncLongContextAttention,
    LongContextAttention,
    set_seq_parallel_pg,
    EXTRACT_FUNC_DICT
)
from yunchang.kernels import AttnType

sp_ulysses_degree = 2
sp_ring_degree = 4

# 支持 world_size = 8
set_seq_parallel_pg(sp_ulysses_degree, sp_ring_degree, rank, world_size)

# attn_type 可以是 FA、FA3 或 TORCH_EFFICIENT。
longctx_attn = LongContextAttention(ring_impl_type="zigzag", attn_type=AttnType.FA)

# 如果您使用的是不支持 flash_attn 的 NPU，可以使用以下代码。
# LongContextAttention(ring_impl_type="zigzag", attn_type=AttnType.TORCH_EFFICIENT)

# 提取全局 Q、K、V 的局部片断。
local_q = EXTRACT_FUNC_DICT["zigzag"](
        Q, rank, world_size=world_size, rd=sp_ring_degree, ud=sp_ulysses_degree
    ).detach().clone()
...

local_out = usp_attn(
        local_q,
        local_k,
        local_v,
        dropout_p=dropout_p,
        causal=True, # zigzag 和 stripe 是针对 causal=True 的负载均衡策略
        window_size=window_size,
        softcap=0.0,
        alibi_slopes=alibi_slopes,
        deterministic=deterministic,
        return_attn_probs=True,
    )

3. 测试

如果您尚未安装 yuanchang，请将项目根目录添加到 PYTHONPATH：

export PYTHONPATH=$PWD:$PYTHONPATH

• FlashAttn/Torch 测试

torchrun --nproc_per_node=4 ./test/test_hybrid_attn.py --sp_ulysses_degree 2 --ring_impl_type "zigzag" --causal --attn_impl fa --use_bwd
torchrun --nproc_per_node=4 ./test/test_hybrid_attn.py --sp_ulysses_degree 2 --ring_impl_type "zigzag" --causal --attn_impl torch
torchrun --nproc_per_node 8 test/test_hybrid_qkvpacked_attn.py

• Sage/SpargeAttention 测试（仅前向传播，非因果）

您需要从源码安装 SpargeAttn 和 SageAttention。

torchrun --nproc_per_node=4 ./test/test_hybrid_attn.py --sp_ulysses_degree 2 --attn_impl sage_fp8

torchrun --nproc_per_node=4 ./test/test_hybrid_attn.py --sp_ulysses_degree 4 --attn_impl sparse_sage --sparse_sage_tune_mode

• FlashInfer 测试（仅前向传播）

请从 这里 安装 FlashInfer。

torchrun --nproc_per_node=4 --master_port=1234 ./test/test_hybrid_attn.py --sp_ulysses_degree 2 --ring_impl_type 'basic_flashinfer' --attn_impl flashinfer

4. 在 Megatron-LM 中验证

数据并行（DP）和统一序列并行（ulysses=2+ring=2）的损失曲线非常接近，如图所示。这一对齐结果证实了统一序列并行的准确性。

当使用负载均衡环形注意力机制并施加因果掩码时，必须使用 EXTRACT_FUNC_DICT 函数对 Query 张量进行重新排序。

在 Megatron-LM 中，您可以在将输入标记送入模型之前对其进行重新排序，并对 RoPE 参数应用相同的重新排序。有关详细说明，请参阅我们的论文。

有关示例实现，您可以查看此 PR，其中将 USP 集成到了 BAAI 的 Megatron-LM 框架中。

6. 基准测试

bash ./scripts/run_qkvpack_compare.sh

在一台 8xA100 NVLink 机器上，基准测试结果如下：

在一台 8xL20 PCIe 机器和一台 4xA100 PCIe 机器上，基准测试结果如下：

一些结论：

1. 如果头数足够多，Ulysses 的性能优于 Ring-Attention。Ulysses 的 All-to-All 通信在单机内非常高效，开销比例极低。相比之下，Ring 将计算和通信分开，这增加了总的计算时间，即使完全重叠，其速度也仍慢于 Ulysses。

2. QKV 打包版本（LongContextAttentionQKVPacked）比 QKV 非打包版本（LongContextAttention）更好，且随着序列长度的缩短，这种差异更加明显。MAQ 和 GQA 只能使用非打包版本。

3. 在 Ring-Attention 的不同实现变体中，zigzag 和 stripe 的性能优于 basic。通常情况下，zigzag 略好于 stripe，但随着序列长度的增加，zigzag 和 stripe 之间的差异变得不那么明显。值得注意的是，zigzag 和 stripe 对序列维度都有特定的布局要求。

4. 混合并行性在异构网络设备上表现良好。例如，在 8-GPU L20 配置中，当 ulysess_degree 设置为 2、ring_degree 设置为 4 时，可达到最佳性能。

7. 4D 并行的最佳实践

我们在一篇技术报告中分析了将序列并行引入数据并行/ZeRO 并行/张量并行/流水线并行的影响，该报告可在 这里 查阅。

以下列出了一些最佳实践：

1. 我们建议使用统一序列并行（Unified-SP）代替 SP-Ring 和 SP-Ulysses，因为它兼具两者的功能，并提供更多优势。

2. 数据并行（DP）与序列并行（SP）：我们建议尽可能优先使用 DP，只有在批量大小（bs）不足以进行划分时，才考虑是否采用 SP。

3. 使用 SP 时，应始终与 ZeRO-1/2 结合使用。

4. 统一序列并行的通信成本低于 Megatron-LM 序列并行下的张量并行！您可以使用统一序列并行来替代张量并行，以获得更好的速度。然而，目前将张量并行切换为 SP+ZeRO2 并不能增加训练中的序列长度。而 SP+ZeRO3 则可以训练与 TP-sp 相似的序列长度。我们建议，在采用 GQA 时，SP 在通信成本方面可能比 TP 更具优势，因为 GQA 可以降低 SP 的通信成本，同时不影响 TP。

5. 可以设置更高的 SP 并行度，这可能需要在头数有限的情况下设置较大的 ring degree，以便在更多计算设备上训练长序列。但 TP 却无法设置较高的并行度。

8. 应用 USP 的项目

我很荣幸本仓库已为以下项目做出了贡献：

1. xdit-project/xDiT
2. NVlabs/VILA
3. feifeibear/Odysseus-Transformer
4. Ascend/AscendSpeed
5. jzhang38/EasyContext
6. FlagOpen/FlagScale
7. zhiyuanhubj/LongRecipe
8. NVIDIA/TransformerEngine
9. xdit-project/mochi-xdit

9. 引用我们

USP：一种用于长上下文生成式 AI 的统一序列并行方法

@article{fang2024unified,
  title={A Unified Sequence Parallelism Approach for Long Context Generative AI},
  author={Fang, Jiarui and Zhao, Shangchun},
  journal={arXiv preprint arXiv:2405.07719},
  year={2024}
}

YunChang (Long-Context-Attention) 快速上手指南

YunChang 是一个统一的序列并行（Unified Sequence Parallelism, USP）注意力机制库，旨在优化长上下文大语言模型（LLM）的训练与推理。它结合了 DeepSpeed-Ulysses 和 Ring-Attention 的优势，解决了两者单独使用时的局限性（如 Ulysses 对头数的限制、Ring-Attention 的计算通信效率问题），特别适用于 GQA/MQA 架构及大规模分布式训练。

1. 环境准备

系统要求

• 硬件：NVIDIA GPU (推荐 A100, H100, L20, L40 等) 或 国产 NPU (如 Ascend)。
• 软件：PyTorch, CUDA (针对 NVIDIA)。

前置依赖

核心依赖为 flash-attn。根据硬件型号选择版本：

• H100/B100 (支持 FA v3)：需从源码安装 FlashAttention v3 (Beta 版)。
• A100/L40 (支持 FA v2)：可直接通过 pip 安装 FlashAttention v2。
• NPU/无 FA 支持硬件：无需安装 flash_attn，将使用 PyTorch 原生实现（注意：原生实现目前不支持反向传播）。

国内加速建议：安装 Python 包时，推荐使用清华或阿里镜像源以提升下载速度。

export PIP_INDEX_URL=https://pypi.tuna.tsinghua.edu.cn/simple

2. 安装步骤

方案 A：标准安装 (适用于大多数 NVIDIA GPU)

直接通过 pip 安装 flash-attn 和 yunchang。

# 安装 flash-attn (自动匹配适合当前硬件的版本，通常为 v2)
pip install flash-attn

# 安装 yunchang
pip install yunchang

方案 B：安装 FlashAttention v3 (仅适用于 H100/B100 等新一代显卡)

由于 FA v3 处于 Beta 阶段，需从源码编译安装。

# 克隆 flash-attention 仓库并切换到指定 commit (2024-11-10)
git clone https://github.com/Dao-AILab/flash-attention.git
cd flash-attention
git checkout b443207c1fc4c98e4532aad4e88cfee1d590d996

# 从源码安装
pip install .

# 返回上级目录并安装 yunchang
cd ..
pip install yunchang

方案 C：本地源码安装 (开发模式)

如果你克隆了本仓库进行二次开发：

pip install .

方案 D：AMD GPU 支持

请参考项目文档 docs/install_amd.md 进行特定安装。

3. 基本使用

以下示例展示了如何使用 zigzag 模式的混合序列并行注意力机制。主要流程包括：设置并行组、提取局部张量、调用注意力计算。

代码示例

import torch
from yunchang import (
    LongContextAttention,
    set_seq_parallel_pg,
    EXTRACT_FUNC_DICT
)
from yunchang.kernels import AttnType

# 1. 配置并行度
# sp_ulysses_degree: Ulysses 并行度
# sp_ring_degree: Ring 并行度
# 确保 rank 和 world_size 与实际运行环境一致 (例如 torchrun 启动的参数)
sp_ulysses_degree = 2
sp_ring_degree = 4
rank = int(os.environ.get("RANK", 0))
world_size = int(os.environ.get("WORLD_SIZE", 8))

# 初始化序列并行进程组
set_seq_parallel_pg(sp_ulysses_degree, sp_ring_degree, rank, world_size)

# 2. 初始化注意力模块
# ring_impl_type: 可选 "zigzag", "stripe", "basic" 等
# attn_type: 可选 AttnType.FA (FlashAttn v2), AttnType.FA3 (FlashAttn v3), AttnType.TORCH_EFFICIENT (无 FA 时)
longctx_attn = LongContextAttention(ring_impl_type="zigzag", attn_type=AttnType.FA)

# 如果是 NPU 或不支持 FlashAttn 的环境，请使用:
# longctx_attn = LongContextAttention(ring_impl_type="zigzag", attn_type=AttnType.TORCH_EFFICIENT)

# 3. 准备数据 (假设 Q, K, V 已定义且形状正确)
# Q, K, V shape: [batch, seq_len, num_heads, head_dim]

# 提取局部切片 (Load Balance 关键步骤)
# 对于 causal mask，必须使用 extract_func 重排输入以平衡负载
local_q = EXTRACT_FUNC_DICT["zigzag"](
    Q, rank, world_size=world_size, rd=sp_ring_degree, ud=sp_ulysses_degree
).detach().clone()

local_k = EXTRACT_FUNC_DICT["zigzag"](
    K, rank, world_size=world_size, rd=sp_ring_degree, ud=sp_ulysses_degree
).detach().clone()

local_v = EXTRACT_FUNC_DICT["zigzag"](
    V, rank, world_size=world_size, rd=sp_ring_degree, ud=sp_ulysses_degree
).detach().clone()

# 4. 执行注意力计算
local_out = longctx_attn(
    local_q,
    local_k,
    local_v,
    dropout_p=0.0,
    causal=True,  # zigzag 和 stripe 策略专为 causal=True 设计以实现负载均衡
    window_size=-1, # -1 表示无滑动窗口限制
    softcap=0.0,
    alibi_slopes=None,
    deterministic=False,
    return_attn_probs=False,
)

# local_out 即为计算结果，后续可根据需要聚合或继续使用

运行测试

安装完成后，可通过以下命令验证安装是否成功（以 4 卡为例）：

# 测试 FlashAttn 实现 (含反向传播)
torchrun --nproc_per_node=4 ./test/test_hybrid_attn.py --sp_ulysses_degree 2 --ring_impl_type "zigzag" --causal --attn_impl fa --use_bwd

# 测试 PyTorch 原生实现 (适用于无 FA 环境，仅前向)
torchrun --nproc_per_node=4 ./test/test_hybrid_attn.py --sp_ulysses_degree 2 --ring_impl_type "zigzag" --causal --attn_impl torch

最佳实践提示

• 并行策略选择：优先使用统一序列并行 (Unified-SP) 替代单独的 Ring-SP 或 Ulysses-SP。
• 与 DP/ZeRO 配合：建议在使用 SP 时同时开启 ZeRO-1 或 ZeRO-2。当显存不足以支撑大 Batch Size 时，再考虑引入 SP。
• GQA/MQA 支持：USP 完美支持 GQA 和 MQA，这是传统 Ulysses 难以做到的。
• 异构网络：在异构网络环境下（如 PCIe 互联），调整 ulysses_degree 和 ring_degree 的比例（例如 2:4）往往能获得更佳性能。