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. |
| 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.
should this be local_rank == 0? What if there isn't a shared filesystem
There was a problem hiding this comment.
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.
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.
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.
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.
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?