Enable universal checkpoint for zero stage 1 #4516
Conversation
tjruwase
requested review from
jeffra,
samyam,
ShadenSmith,
duli2012,
RezaYazdaniAminabadi,
awan-10,
cmikeh2 and
arashb
as code owners
tjruwase
removed request for
arashb,
ShadenSmith,
cmikeh2,
duli2012,
samyam,
awan-10 and
RezaYazdaniAminabadi
|
@stas00, FYI |
tjruwase
changed the title
|
Amazing! Thank you for starting to work on this super-essential feature, Tunji! |
…d logging circular import issue
…o olruwase/ds_2921
* Enable uni_ckpt for z1 * Remove logging fix to seperate PR. Relocate conversion script to avoid logging circular import issue * Formatting fix * PR feedback * Handle replicated params * Detect bf16_optimizer * Docs * Fix docs
* Enable uni_ckpt for z1 * Remove logging fix to seperate PR. Relocate conversion script to avoid logging circular import issue * Formatting fix * PR feedback * Handle replicated params * Detect bf16_optimizer * Docs * Fix docs
| os.makedirs(param_base_path, exist_ok=True) | ||
|
|
||
| cnt += 1 | ||
| counter = f"{dp_index:0>2d}" |
There was a problem hiding this comment.
Dear @tjruwase ,
I'm currently examining a scenario where the maximum dp_index is 127, which, according to alphabetical order, is considered smaller than 13. This raises a question regarding the tensor sorting process in line 144 of our code. Given this context, could there potentially be any issues with how the tensors are sorted due to this ordering?
I appreciate your insight on this matter.
Best regards,
Junfeng
There was a problem hiding this comment.
A specific example is as follows:
./temp/model.transformer_encoder.layers.19.self_attn.in_proj_weight => fp32.100, torch.Size([2187604])
./temp/model.transformer_encoder.layers.19.self_attn.in_proj_weight => fp32.99, torch.Size([48144044])
There was a problem hiding this comment.
@rgtjf, thanks for sharing this issue. Do you mind creating a new ticket for it? I can see that line 130 may not generalize to larger scales. It would be great if you could share more details in a new ticket. Thanks!
Generalize universal checkpointing in DS:
Tested with Megatron-DS GPT using companion PR microsoft/Megatron-DeepSpeed#265
Fix #2921
[ ] Tutorial: Defer