Adding the new feature of FPDT #441
Open
+50,600
−113
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
…on for supporting batch size larger than 1
YJHMITWEB
requested review from
tjruwase,
awan-10,
eltonzheng,
duli2012,
arashb and
GuanhuaWang
as code owners
|
Hi @YJHMITWEB , is FPDT referring to this paper? https://ui.adsabs.harvard.edu/abs/2023JARS...17b6510H/abstract |
|
@YJHMITWEB Do we need changes in |
tohtana
reviewed
Aug 30, 2024
megatron/initialize.py
Outdated
| @@ -349,9 +349,12 @@ def _warmup_jit_function(): | |||
| dtype = torch.float32 | |||
|
|
|||
| # Warmup fused bias+gelu | |||
| seq_length = args.seq_length | |||
| if args.ds_sequence_parallel_fpdt: | |||
| seq_length = 8192 | |||
There was a problem hiding this comment.
Can you define this as another variable like "FPDT_SEQ_LEN" and give a description in a comment why we have this setting?
There was a problem hiding this comment.
This is fixed by setting it to be ds_sequence_parallel_fpdt_chunk_size if FPDT is enabled.
inkcherry
reviewed
Aug 30, 2024
| @@ -32,7 +35,9 @@ def forward(self, max_seq_len, offset=0): | |||
| emb = torch.cat((freqs, freqs), dim=-1) | |||
| # emb [seq_length, .., dim] | |||
| from einops import rearrange | |||
| return rearrange(emb, 'n d -> n 1 1 d') | |||
| base = rearrange(emb, 'n d -> n 1 1 d') | |||
There was a problem hiding this comment.
will this change the output when use --use-rotary-position-embeddings, in llama style model?
FYI https://github.com/microsoft/Megatron-DeepSpeed/blob/main/examples_deepspeed/pretrain_llama2_distributed.sh
There was a problem hiding this comment.
We have tested both GPT and Llama models, this works well with both.
@delock, no, FPDT refers to this paper, aka Ulysses-Offload |
Thanks @samadejacobs for pointing. |
|
|
@microsoft-github-policy-service agree |
github-merge-queue bot
pushed a commit
to microsoft/DeepSpeed
that referenced
this pull request
Nov 26, 2024
[FPDT](https://arxiv.org/abs/2408.16978) can only be used with [this version](microsoft/Megatron-DeepSpeed#441) of Megatron-DeepSpeed. --------- Co-authored-by: Jinghan Yao <[email protected]> Co-authored-by: Sam Ade Jacobs <[email protected]> Co-authored-by: Jinghan Yao <[email protected]> Co-authored-by: Logan Adams <[email protected]> Co-authored-by: Olatunji Ruwase <[email protected]> Co-authored-by: Jinghan Yao <[email protected]> Co-authored-by: Logan Adams <[email protected]> Co-authored-by: Masahiro Tanaka <[email protected]> Co-authored-by: Masahiro Tanaka <[email protected]>
github-merge-queue bot
pushed a commit
to microsoft/DeepSpeed
that referenced
this pull request
Nov 26, 2024
[FPDT](https://arxiv.org/abs/2408.16978) can only be used with [this version](microsoft/Megatron-DeepSpeed#441) of Megatron-DeepSpeed. --------- Co-authored-by: Jinghan Yao <[email protected]> Co-authored-by: Sam Ade Jacobs <[email protected]> Co-authored-by: Jinghan Yao <[email protected]> Co-authored-by: Logan Adams <[email protected]> Co-authored-by: Olatunji Ruwase <[email protected]> Co-authored-by: Jinghan Yao <[email protected]> Co-authored-by: Logan Adams <[email protected]> Co-authored-by: Masahiro Tanaka <[email protected]> Co-authored-by: Masahiro Tanaka <[email protected]>
github-merge-queue bot
pushed a commit
to microsoft/DeepSpeed
that referenced
this pull request
Nov 26, 2024
[FPDT](https://arxiv.org/abs/2408.16978) can only be used with [this version](microsoft/Megatron-DeepSpeed#441) of Megatron-DeepSpeed. --------- Co-authored-by: Jinghan Yao <[email protected]> Co-authored-by: Sam Ade Jacobs <[email protected]> Co-authored-by: Jinghan Yao <[email protected]> Co-authored-by: Logan Adams <[email protected]> Co-authored-by: Olatunji Ruwase <[email protected]> Co-authored-by: Jinghan Yao <[email protected]> Co-authored-by: Logan Adams <[email protected]> Co-authored-by: Masahiro Tanaka <[email protected]> Co-authored-by: Masahiro Tanaka <[email protected]>
github-merge-queue bot
pushed a commit
to microsoft/DeepSpeed
that referenced
this pull request
Nov 26, 2024
[FPDT](https://arxiv.org/abs/2408.16978) can only be used with [this version](microsoft/Megatron-DeepSpeed#441) of Megatron-DeepSpeed. --------- Co-authored-by: Jinghan Yao <[email protected]> Co-authored-by: Sam Ade Jacobs <[email protected]> Co-authored-by: Jinghan Yao <[email protected]> Co-authored-by: Logan Adams <[email protected]> Co-authored-by: Olatunji Ruwase <[email protected]> Co-authored-by: Jinghan Yao <[email protected]> Co-authored-by: Logan Adams <[email protected]> Co-authored-by: Masahiro Tanaka <[email protected]> Co-authored-by: Masahiro Tanaka <[email protected]>
github-merge-queue bot
pushed a commit
to microsoft/DeepSpeed
that referenced
this pull request
Dec 2, 2024
[FPDT](https://arxiv.org/abs/2408.16978) can only be used with [this version](microsoft/Megatron-DeepSpeed#441) of Megatron-DeepSpeed. --------- Co-authored-by: Jinghan Yao <[email protected]> Co-authored-by: Sam Ade Jacobs <[email protected]> Co-authored-by: Jinghan Yao <[email protected]> Co-authored-by: Logan Adams <[email protected]> Co-authored-by: Olatunji Ruwase <[email protected]> Co-authored-by: Jinghan Yao <[email protected]> Co-authored-by: Logan Adams <[email protected]> Co-authored-by: Masahiro Tanaka <[email protected]> Co-authored-by: Masahiro Tanaka <[email protected]>
samadejacobs
reviewed
Dec 2, 2024
| key_layer = mixed_x_layer[:, :, self.projection_size:self.projection_size+self.kv_projection_size].reshape(seq_len, bs, -1, self.head_dim) | ||
| value_layer = mixed_x_layer[:, :, self.projection_size+self.kv_projection_size:].reshape(seq_len, bs, -1, self.head_dim) | ||
| if self.sequence_parallel or not self.enable_ds_sequence_parallel: | ||
| seq_len, bs = mixed_x_layer.shape[0], mixed_x_layer.shape[1] |
There was a problem hiding this comment.
Might be worthwhile to keep "split_tensor" implementation as default, that is not fpdt scenario.
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
FPDT can only be work with this version of DeepSpeed.