implement mid-train-chkpting#791
Conversation
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 8be76d87e9
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| if should_step or is_last: | ||
| self._maybe_save_mid_train_checkpoint( | ||
| epoch=epoch, | ||
| batch_in_epoch=global_data_iter_step, | ||
| ) |
There was a problem hiding this comment.
When the interval expires on the final training batch, this branch writes ckpt_mid_train.pt with epoch_complete=False and batch_in_epoch set to the last batch. If the job is preempted during validation/inference before ckpt.pt is written, load_checkpoint() resumes the same epoch at last_batch + 1, _OffsetBatchSampler yields no batches, and the train aggregator has no valid batch data for train/mean/loss. Skip mid-train saves when is_last is true, or mark/save this state as an epoch-complete checkpoint after scheduler/epoch bookkeeping is complete.
Useful? React with 👍 / 👎.
| batch_sampler = _OffsetBatchSampler( | ||
| raw_loader.batch_sampler, | ||
| start_batch=start_batch_in_epoch, | ||
| ) |
There was a problem hiding this comment.
For single-process CPU/GPU runs, the wrapped raw_loader.batch_sampler is the non-distributed shuffled sampler, which creates a fresh random batch order when iterated. Since the checkpoint only stores the ordinal batch_in_epoch, resuming from a later epoch after set_seed() resets RNG skips that many batches in a new order, silently replaying some already-trained batches and skipping some untrained ones. Make this sampler order deterministic by epoch and resume from that same order, or save/restore the sampler RNG/order with the checkpoint.
Useful? React with 👍 / 👎.
| except Exception: | ||
| logger.exception("Failed to save mid-train checkpoint to %s", checkpoint_path) | ||
| return None |
There was a problem hiding this comment.
If writing ckpt_mid_train.pt fails here (for example due to disk full, permissions, or a transient filesystem error), training only logs the exception and continues, leaving preemptible runs without the checkpoint they were configured to rely on. That also contradicts the root AGENTS.md guidance to not swallow errors; please re-raise after logging so the failure is visible like the existing epoch checkpoint saves.
Useful? React with 👍 / 👎.
1 participant
implement mid-train-checkpointing with default=off. If set to n minutes, it saves a mid-epoch checkpoint every n minutes. This enables periodic saving within epochs to prevent a crash or training termination from resulting in significant loss of training effort. This is most useful when epochs take a very long time to train.