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:
2026-07-13 12:27:09 +08:00
parent 9bcd696580
commit b092316385
3 changed files with 42 additions and 22 deletions
-3
View File
@@ -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,