add checkpoint #945
add checkpoint #945
Conversation
|
Looks good to me. cc @tjruwase |
@hwchen2017 , this is a good question, if this is standard in pytorch or megatron, we should keep, otherwise we can skip it. |
| @@ -16,7 +16,7 @@ | |||
| from megatron import get_tensorboard_writer | |||
| from megatron.core import mpu, tensor_parallel | |||
| from megatron.arguments import parse_args, validate_args | |||
| from megatron.checkpointing import load_args_from_checkpoint | |||
| # from megatron.checkpointing import load_args_from_checkpoint | |||
There was a problem hiding this comment.
Delete instead of comment.
| @@ -241,7 +247,7 @@ def save_checkpoint(iteration, model, optimizer, opt_param_scheduler): | |||
|
|
|||
| # Arguments, iteration, and model. | |||
| state_dict = {} | |||
| state_dict['args'] = args | |||
| # state_dict['args'] = args | |||
There was a problem hiding this comment.
It is better to preserve args as it stores useful information for restarts.
There was a problem hiding this comment.
@tjruwase thanks for the suggestion. I have fixed it.
This is a reason why |
|
@zhangsmallshark please address above comments. Thanks! |
|
@GuanhuaWang I fixed it. Please check it. |
|
@zhangsmallshark - could you sign off with DCO on this PR? It replaces the CLA we had before. To fix it, the steps should be here |
support checkpoint for domino