Add RingFlashAttention for context parallel #8383
Conversation
|
Thanks for your contribution! |
|
|
zhangyuqin1998
force-pushed
the
ring_flash_attention
branch
from
09fd62e
to
fbd16a1
Compare
Codecov ReportAttention: Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## develop #8383 +/- ##
===========================================
- Coverage 54.22% 53.87% -0.35%
===========================================
Files 617 620 +3
Lines 96203 97065 +862
===========================================
+ Hits 52164 52297 +133
- Misses 44039 44768 +729 ☔ View full report in Codecov by Sentry. |
|
|
||
| # if step != cp_size - 1: | ||
| # comm_buffer.wait() | ||
| paddle.device.synchronize() |
There was a problem hiding this comment.
TODO：batch_isend_irecv异步流下，无法wait，需要修复。对性能有影响。
There was a problem hiding this comment.
TODO：batch_isend_irecv异步流下，无法wait，需要修复。对性能有影响。
done～
zhangyuqin1998
force-pushed
the
ring_flash_attention
branch
from
f94a915
to
4e88520
Compare
zhangyuqin1998
force-pushed
the
ring_flash_attention
branch
2 times, most recently
from
cf7d334
to
88bc460
Compare
| block_out, _, block_lse, _ = _C_ops.flash_attn( | ||
| local_query, | ||
| block_k[:, : local_q_seq_len // 2, :, :], | ||
| block_v[:, : local_q_seq_len // 2, :, :], |
There was a problem hiding this comment.
这种方式，性能可能比较慢。看看能否直接使用op的方式调用。
There was a problem hiding this comment.
这种方式，性能可能比较慢。看看能否直接使用op的方式调用。
done～
| if attn_mask is not None: | ||
| attn_masks_list = paddle.split(attn_mask, num_or_sections=cp_size * 2, axis=3) | ||
| if is_causal: | ||
| local_query_second_chunk = local_query[:, local_q_seq_len // 2 :, :, :].clone().contiguous() |
There was a problem hiding this comment.
contiguous ？ 可能不需要这个。尽量使用切分的api，不实用运算符重载。
There was a problem hiding this comment.
contiguous ？ 可能不需要这个。尽量使用切分的api，不实用运算符重载。
done～
| grad_comm_buffer = RingCommunicator(group, key_grad_buffer, value_grad_buffer) | ||
|
|
||
| if is_causal: | ||
| local_query_second_chunk = local_query[:, local_q_seq_len // 2 :, :, :].clone().contiguous() |
There was a problem hiding this comment.
这个前向已经计算过了，是否可以优化不计算。
已做优化
| def wait(self): | ||
| # for req in self._reqs: | ||
| # req.wait() | ||
| # self._reqs = None |
There was a problem hiding this comment.
改成TODO吧。不用注释。
done～
zhangyuqin1998
force-pushed
the
ring_flash_attention
branch
from
a360468
to
ab562b7
Compare
paddlenlp/trainer/training_args.py
Outdated
| @@ -583,6 +587,15 @@ class TrainingArguments: | |||
| ) | |||
| }, | |||
| ) | |||
| cp_parallel_degree: int = field( | |||
There was a problem hiding this comment.
换成 context_parallel_degree ？
| cp_parallel_degree: int = field( | |
| context_parallel_degree: int = field( |
paddlenlp/trainer/trainer.py
Outdated
| @@ -763,6 +764,8 @@ def train( | |||
| trainable_numel = int(trainable_numel_tensor.item()) // self.args.dataset_world_size | |||
| if self.args.sep_parallel_degree > 0: | |||
| trainable_numel = trainable_numel // self.args.sep_parallel_degree | |||
| if self.args.cp_parallel_degree > 0: | |||
| trainable_numel = trainable_numel // self.args.cp_parallel_degree | |||
paddlenlp/trainer/training_args.py
Outdated
| @@ -230,6 +230,10 @@ class TrainingArguments: | |||
| The paddle sequence parallel strategy. It can reduce the GPU memory of activation to 1/sep, and it is orthogonal to | |||
| data parallel, sharding stage1, tensor parallel and pipeline parallel strategy. | |||
| ) | |||
| cp_parallel_degree (`int`, *optional*, defaults to `-1`)( | |||
There was a problem hiding this comment.
这个参数在 docs/trainer.md 文档中也加一下吧。
| self.tensor_parallel_degree | ||
| * self.sep_parallel_degree | ||
| * self.cp_parallel_degree | ||
| * self.pipeline_parallel_degree |
PR types
New features
PR changes
Models
Description
为fleet的context parallel增加ring flash attention的支持
paddle兼容性:
使用paddle中的sep group，对paddle无改动
收敛性：
将cp和sep做对比。理论上，二者的收敛结果应该完全一致。经过测试，sep和cp的收敛情况近乎一致。绿色为cp，蓝色为sep。
性能：
单机8卡小模型测试，序列长度为20k时，性能对比如图。绿色为cp，蓝色为sep。