关于数据分割和合并

[Kwen-Chen]

您好，我注意到代码里有这样的一部分

    local_q = q.chunk(world_size, dim=1)[rank].detach().clone()
    local_q.requires_grad = True
    local_k = k.chunk(world_size, dim=1)[rank].detach().clone()
    local_k.requires_grad = True
    local_v = v.chunk(world_size, dim=1)[rank].detach().clone()
    local_v.requires_grad = True

    local_dout = dout.chunk(world_size, dim=1)[rank].detach().clone()

在调用 dist_attention 时，需要先将数据分割吗？那么这样算出来的 attention_output，该如何合并呢？

[feifeibear]

这里是模拟序列并行，每个GPU存1/N的Q、K、V矩阵。输出的attn output就是(bs, seqlen/N, hidden_size)，你可以直接给后面的FNN做张量并行、ZeRO等其他并行。

你所谓的合并（把attn output变成多卡replica的？）在实际场景不存在。

[Kwen-Chen]

这里是模拟序列并行，每个GPU存1/N的Q、K、V矩阵。输出的attn output就是(bs, seqlen/N, hidden_size)，你可以直接给后面的FNN做张量并行、ZeRO等其他并行。

你所谓的合并（把attn output变成多卡replica的？）在实际场景不存在。

十分感谢您的回复，也就是说，对于 dist_attn，在 forward 时，送入的 K、Q、V 应该是被分割 seqlen/N 后的 K、Q、V？那么后续在使用 attn output 时，都需要使用 (bs, seqlen/N, hidden_size) 的吗？ 因为我在做 长序列训练 LLama 的实验，我本以为只需要改变 Attention （如果教程中的一样），这是否意味着，还需要改变 Llama 中后续对 attn output 的处理，改成 (bs, seqlen/N, hidden_size) ？

[Kwen-Chen]

另外我注意到，您更改了 原 L26 中的 AlltoAll 函数，在 L26 里是否已做了 K、Q、V的分割？

[feifeibear]

根据QKV是否packed在一起，有两个版本。 LongContextAttentionQKVPacked，LongContextAttention。

没有看deepspeed的all2all，它输入要求比较奇怪，本项目重新实现了。

[Kwen-Chen]

抱歉，您理解错了我的意思，我看了一下午，但始终没有搞明白，还请赐教：

1. 是否意思为在模型训练的时候使用 N/P 长度的数据进入训练，这样用 （N/P，d） 的Q、K、V 算出来的 Output 也是（ N/P，d ）的，用这个 Output 去做 backward。但疑问在于，这样的训练方式难道各自不是独立的吗？context length 应该是 N/P，而不是N啊
2. 所以原本的思路是，各个卡跑的是 N/P 长度的数据，但最终汇总为一个 Output （N，d），这样再做 backward，context length才是N吧。

我对这里可能一直有理解误区，实在抱歉，还请赐教

[ShomyLiu]

@TheCoderChen 我说下我的理解，以ulysses为例，可以看下ulysses的论文的图， 前面有分割QKV，但是在计算注意力之前，使用 All-To-All 通信把QKV又聚合起来了，类似转置操作，保证每个卡上计算的是完整长度了，只不过每个卡计算的不是所有的计算头，而是一部分，最后计算完之后，再结合AlltoAll，得到所有头的结果。 不知道解释的是否有问题 @feifeibear

[Kwen-Chen]

@ShomyLiu 感谢你的回复，是这样的，我之前的理解有些浅显，这里即使是传入的 （N/P，d） 的Q、K、V ，通过 All-To-All 之后，也绝不是彼此独立的，是每张卡上计算的是完整长度。那么得到所有头之后的结果后该如何处理呢？如果做训练的话，直接送进去计算梯度后对每个 N/P 计算的梯度直接进行 Allreduce 是否正确呢？我主要对这一块不太理解

[feifeibear]

十分感谢您的回复，也就是说，对于 dist_attn，在 forward 时，送入的 K、Q、V 应该是被分割 seqlen/N 后的 K、Q、V？那么后续在使用 attn output 时，都需要使用 (bs, seqlen/N, hidden_size) 的吗？ 因为我在做 长序列训练 LLama 的实验，我本以为只需要改变 Attention （如果教程中的一样），这是否意味着，还需要改变 Llama 中后续对 attn output 的处理，改成 (bs, seqlen/N, hidden_siz

Q：是否意思为在模型训练的时候使用 N/P 长度的数据进入训练，这样用 （N/P，d） 的Q、K、V 算出来的 Output 也是（ N/P，d ）的，用这个 Output 去做 backward。但疑问在于，这样的训练方式难道各自不是独立的吗？context length 应该是 N/P，而不是N啊 A：第一个句号前是对的。这样的训练方式难道各自不是独立的吗。我没理解你的这句话的问题。

Q：所以原本的思路是，各个卡跑的是 N/P 长度的数据，但最终汇总为一个 Output （N，d），这样再做 backward，context length才是N吧。 A：原本思路如果指Ulysses或ring，你理解有误。它们每个GPU输出都是(N/P, d)

[fmo-mt]

Q：是否意思为在模型训练的时候使用 N/P 长度的数据进入训练，这样用 （N/P，d） 的Q、K、V 算出来的 Output 也是（ N/P，d ）的，用这个 Output 去做 backward。但疑问在于，这样的训练方式难道各自不是独立的吗？context length 应该是 N/P，而不是N啊 A：第一个句号前是对的。这样的训练方式难道各自不是独立的吗。我没理解你的这句话的问题。

Q：所以原本的思路是，各个卡跑的是 N/P 长度的数据，但最终汇总为一个 Output （N，d），这样再做 backward，context length才是N吧。 A：原本思路如果指Ulysses或ring，你理解有误。它们每个GPU输出都是(N/P, d)

他说的“各自独立”是想指，如果每个GPU上保持 N/P 的前向+反向，那GPU之间的数据就没有相互依赖，成各自独立训练了。我理解最后的输出肯定是要Gather起来再进行bwd？

[Kwen-Chen]

Q：是否意思为在模型训练的时候使用 N/P 长度的数据进入训练，这样用 （N/P，d） 的Q、K、V 算出来的 Output 也是（ N/P，d ）的，用这个 Output 去做 backward。但疑问在于，这样的训练方式难道各自不是独立的吗？context length 应该是 N/P，而不是N啊 A：第一个句号前是对的。这样的训练方式难道各自不是独立的吗。我没理解你的这句话的问题。 Q：所以原本的思路是，各个卡跑的是 N/P 长度的数据，但最终汇总为一个 Output （N，d），这样再做 backward，context length才是N吧。 A：原本思路如果指Ulysses或ring，你理解有误。它们每个GPU输出都是(N/P, d)

他说的“各自独立”是想指，如果每个GPU上保持 N/P 的前向+反向，那GPU之间的数据就没有相互依赖，成各自独立训练了。我理解最后的输出肯定是要Gather起来再进行bwd？

对，我是你说的这个意思，所以这里最后的输出是要 gather 之后再 bwd吗？这一块有参考实现吗？

[LzhinFdu]

Q：是否意思为在模型训练的时候使用 N/P 长度的数据进入训练，这样用 （N/P，d） 的Q、K、V 算出来的 Output 也是（ N/P，d ）的，用这个 Output 去做 backward。但疑问在于，这样的训练方式难道各自不是独立的吗？context length 应该是 N/P，而不是N啊 A：第一个句号前是对的。这样的训练方式难道各自不是独立的吗。我没理解你的这句话的问题。 Q：所以原本的思路是，各个卡跑的是 N/P 长度的数据，但最终汇总为一个 Output （N，d），这样再做 backward，context length才是N吧。 A：原本思路如果指Ulysses或ring，你理解有误。它们每个GPU输出都是(N/P, d)

他说的“各自独立”是想指，如果每个GPU上保持 N/P 的前向+反向，那GPU之间的数据就没有相互依赖，成各自独立训练了。我理解最后的输出肯定是要Gather起来再进行bwd？

对，我是你说的这个意思，所以这里最后的输出是要 gather 之后再 bwd吗？这一块有参考实现吗？

并行训练时，会在梯度优化之前reduce各个进程上的参数梯度。一般accelerator、deepspeed这些会默认支持。

[feifeibear]

Q：是否意思为在模型训练的时候使用 N/P 长度的数据进入训练，这样用 （N/P，d） 的Q、K、V 算出来的 Output 也是（ N/P，d ）的，用这个 Output 去做 backward。但疑问在于，这样的训练方式难道各自不是独立的吗？context length 应该是 N/P，而不是N啊 A：第一个句号前是对的。这样的训练方式难道各自不是独立的吗。我没理解你的这句话的问题。 Q：所以原本的思路是，各个卡跑的是 N/P 长度的数据，但最终汇总为一个 Output （N，d），这样再做 backward，context length才是N吧。 A：原本思路如果指Ulysses或ring，你理解有误。它们每个GPU输出都是(N/P, d)

他说的“各自独立”是想指，如果每个GPU上保持 N/P 的前向+反向，那GPU之间的数据就没有相互依赖，成各自独立训练了。我理解最后的输出肯定是要Gather起来再进行bwd？

对，我是你说的这个意思，所以这里最后的输出是要 gather 之后再 bwd吗？这一块有参考实现吗？

https://github.com/feifeibear/long-context-attention/blob/main/patches/Megatron-DeepSpeed.patch

你看这个patch