hanwen-sun
commented
1 month ago 具体实现
- 正向公式:
/ cos(i) * x(i) - sin(i) * x(i + Size / 2) i < Size / 2 out(i) = \ cos(i) * x(i) + sin(i) * x(i - Size / 2) i >= Size / 2
其中, out为输出数组，x为输入数组, cos与sin为常量数组; Size为 rotary_size / head_size, 常量;
- 参数解析:
在具体的kernel实现中，有一些比较重要的参数:
B: batch_size, M: seq_len, H: num_heads, K: head_size;
x_layout: BHMK等形式 output_layout: BHMK等形式
k_size = head_size;
tensor_index: 暂时不理解;
还有部分中间参数可以结合代码输出理解;
- 反向公式:
<img width="924" alt="d52a19ba4e15f7e7e877e6ede3d949e" src="https://github.com/Oneflow-Inc/oneflow/assets/83819142/2f8d675d-0d7c-437f-94d1-b66d47d7e536">
<img width="771" alt="7fe69e6a635500fc87503da46961814" src="https://github.com/Oneflow-Inc/oneflow/assets/83819142/d660a285-74c0-4920-ab50-a72d7a381a11">
/ cos(i) * x(i) + sin(i + Size / 2) * x(i + Size / 2) i < Size / 2out(i) =
\ cos(i) x(i) - sin(i - Size / 2) x(i - Size / 2) i >= Size / 2
其中 x为上层传入的反向梯度数组，out为传入下一层的反向梯度数组;
可以对比apex实现，与推导一致: https://github.com/NVIDIA/apex/blob/a7de60e57f0534266841e1733262601ad76aaa74/csrc/megatron/fused_rotary_positional_embedding.h#L175
- 实现细节:
- fused_apply_rotary_emb处理:
- 注意判断saved_tensors的个数，输入tensors的个数可能为1 - 4个:
- 1: x 2: x, position_ids, 3: x, cos, sin 4: s, cos, sin, position_ids, 注: 我们默认cos与sin数组同时出现;
- 为了实现方便，我们将输入的x tensor依旧传入反向算子，实际实现可能可以省略;
- fused_attention_ops.cpp:
- 总体处理逻辑与正向算子保持一致;
- fused_attention_kernel 中 FusedApplyRotaryEmbGradKernel实现:
- 在调用流程中都添加 is_forward判断，在实际的launchkernel中进行分流;
- 添加planegradkernel, 其中cos_val与sin_val的生成不确定实现是否正确;
- 添加intervalgradkernel(未实现)
- test_fused_rotary_embedding
- 实现tensor版本的前向散op, 直接计算反向传播的梯度，与融合算子的反向进行梯度精度对比;
- test_plane需要更详细的测试用例;
内容:
TODO: