-
Notifications
You must be signed in to change notification settings - Fork 984
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Fix checkpoint / config file management #463
Conversation
It’s probably worth testing with MP > 1, PP > 1, and two machines. That’s about as maximally asynchronous as it’ll be. |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
LGTM but I would remove some of the comments for the tag name and just double check the torch.distributed.get_rank() == 0
thing
model.save_checkpoint(neox_args.save, tag=tag, client_state=sd) | ||
|
||
# save config files | ||
if torch.distributed.get_rank() == 0 and neox_args.config_files is not None: |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
should this be local_rank == 0? What if there isn't a shared filesystem
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Following reasoning. Hence I would like to keep it only for rank 0.
- also without shared file system the checkpoint is only complete with data from all ranks. There is no added benefit, you would not use only part of the checkpoint.
- duplicating data across the ranks then would require overwrites when moving the files to one location
- I have had at least one shared file system with problems writing to the same file from multiple ranks. This would avoid such multiple write handles. Deepspeed does it as well for latest file and the zero conversion script. Maybe should not be a problem, but it seems to be one some times.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Deepspeed DOES write to one file at the same time from multiple ranks. To avoid confusion…
@@ -709,6 +709,11 @@ class NeoXArgsTraining(NeoXArgsTemplate): | |||
Output directory to save checkpoints to. | |||
""" | |||
|
|||
config_files: dict = None |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I'm wondering if we should initialize non-configurable variables like this in post_init - I don't think they really warrant a docstring, so it might be a better place for them.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
My thinking was to facilitate development for all not having seen this PR by making all properties of neox args transparent. If some contributor wants to use this info it is nice to read it in the docstring and not be required to read through code to actually see that it exists and what it is.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Ah also: the way it is implemented the file contents are part of neox args and distributed by the deepspeed launcher. This means it needs to be a named argument in the current implementation. So the property needs to be defined and then probably also be documented.
Can do but probably earliest tomorrow. |
Now also tested with pp 2 and mp4 on two nodes with 8 gpus each. Works. |
@sdtblck How does this look to you? It looks ready to merge to me. |
Saves config files in the global step directory instead of the neox_args.save directory. Discussion here: #459
Tested on small config with:
Question is how to get the config files from the load in deepy to the single processes. I opted for the hand-trough via neox args. Any other idea?