Adding the new feature of FPDT #6462
Conversation
deepspeed/sequence/fpdt_layer.py
Outdated
| return out, lse | ||
|
|
||
|
|
||
| def single_all_to_all(input_, scatter_idx, gather_idx, group): |
There was a problem hiding this comment.
Hi @YJHMITWEB is this single_all_to_all the same as the one in
DeepSpeed/deepspeed/sequence/layer.py
Line 41 in 89c4d9f
There was a problem hiding this comment.
We have not tested the non-blocking All2all with our FPDT design, therefore, we use the original non-blocking version. If this is preferred, we can test this.
There was a problem hiding this comment.
This is fixed. We will now use the single_all_to_all from layer.py
deepspeed/sequence/fpdt_layer.py
Outdated
|
|
||
| compute_stream.wait_stream(offload_stream) | ||
| compute_stream.synchronize() | ||
| with torch.cuda.stream(offload_stream): |
There was a problem hiding this comment.
torch.cuda.stream(stream) should be replaced by get_accelerator().stream(stream), same of other occurance in this file.
deepspeed/sequence/fpdt_layer.py
Outdated
| grad_qkv_linear_bias = torch.zeros(qkv_linear_bias.shape, device=qkv_linear_weight.device, dtype=torch.float) | ||
|
|
||
| grad_global_attn_output_chunk = single_all_to_all(grad_output[:, :chunk_size].contiguous(), scatter_idx, gather_idx, spg) | ||
| torch.cuda.synchronize() |
There was a problem hiding this comment.
Here is should be get_accelerator().synchronize()
There was a problem hiding this comment.
This one is solved.
deepspeed/sequence/layer.py
Outdated
| @@ -168,7 +196,8 @@ def __init__( | |||
| self.gather_idx = gather_idx | |||
| self.sp_overlap_comm = False | |||
| self.overlap_handles = None | |||
| self.sp_stream = sp_stream | |||
| self.sp_stream = None | |||
| # self.sp_stream = sp_stream | |||
There was a problem hiding this comment.
Why comment out this line?
There was a problem hiding this comment.
This one is solved.
fix format and add unit test for fpdt
| - '.github/workflows/nv-flash-attn.yml' | ||
| concurrency: | ||
| group: ${{ github.workflow }}-${{ github.ref }} | ||
| cancel-in-progress: true |
There was a problem hiding this comment.
Should we add a cron trigger as well to this? And perhaps the Nightly CI failure issue creation from here - that way if a PR doesn't trigger this we will still know if it fails?
There was a problem hiding this comment.
Thank you @loadams for the suggestion! I added both. Can you check them?
FPDT can only be used with this version of Megatron-DeepSpeed.