feat : add distributed checkpoint via executor checkpoint_context
- Add checkpoint_context context manager to BaseExecutor with entry/exit barrier - Add _gather_state_dict hook overridden per executor (template method) - DDPExecutor skips unwrap on non-rank-0 to avoid redundant state_dict gather - FSDPExecutor uses rank0_only=True to reduce memory on non-writers - Remove redundant rank-0 guard from Checkpoint.save and manual barrier from Callback
This commit is contained in:
@@ -148,9 +148,6 @@ class Checkpoint:
|
||||
save_path = Path(save_dir)
|
||||
save_path.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
if get_rank() != 0:
|
||||
return
|
||||
|
||||
meta = {
|
||||
"epoch": self.epoch,
|
||||
"consumed_samples": self.consumed_samples,
|
||||
|
||||
Reference in New Issue
Block a user