"Save best" and other improvements to model checkpointing, Issue #307

[surojitiitg]

Summary

This PR improves checkpoint saving behavior during training by adding configurable checkpoint frequency and stable best-checkpoint tracking.

Changes

• Files modified: src/speculators/train/trainer.py, src/speculators/train/checkpointer.py, scripts/train.py
• Files added: tests/unit/train/test_checkpoint.py
• Added --checkpoint-freq to control how often numbered checkpoints are saved
• Added --save-best to maintain a stable checkpoint_best/ path
• Implemented symlink support so checkpoint_best/ points to the best saved checkpoint directory
• Kept default behavior backward compatible by using checkpoint_freq=1
• Updated validation flow in src/speculators/train/trainer.py so that val_epoch(...) returns the computed val_metrics dictionary to the training loop
• Best-checkpoint selection uses the key loss_epoch from the dictionary returned by self.val_epoch(epoch) as the monitored metric. The implementation is carried out in src/speculators/train/trainer.py
• Added tests/unit/train/test_checkpoint.py, which contains unit tests for:
  • ignoring checkpoint_best during resume checkpoint discovery
  • creating and updating the best-checkpoint symlink
  • checkpoint selection behavior during training

Behavior

• If --checkpoint-freq and --save-best are not provided, the current checkpointing behavior is preserved (backward compatible defaults).
• checkpoint_freq=1 preserves the current behavior of saving checkpoints every epoch
• When save_best=True, checkpoint_best/ points to the checkpoint with the lowest validation loss among the saved checkpoints
• Best-checkpoint tracking is based on val_metrics["loss_epoch"] returned from the function val_epoch(...)
• Resume behavior still loads the latest numbered checkpoint and ignores checkpoint_best/

Notes

• The first epoch checkpoint is intentionally saved as part of the current behavior. This provides a a guaranteed recovery point and early verification artifact.
• checkpoint_best/ is implemented as a symlink to avoid duplicating checkpoint files

Testing

Tested with:

• Unit tests for checkpoint selection and symlink behavior. The test file is located in tests/unit/train/test_checkpoint.py
• Did testing to observe the file structure of the checkpointing via scripts/train.py using:

CUDA_VISIBLE_DEVICES=1 python scripts/train.py \
  --verifier-name-or-path meta-llama/Llama-3.1-8B-Instruct \
  --data-path output \
  --save-path checkpoints_smoke \
  --log-dir logs_smoke \
  --epochs 2 \
  --lr 3e-5 \
  --total-seq-len 256 \
  --data-format-version 1 \
  --num-layers 1 \
  --d2t-path vocab_mapping/d2t.npy \
  --t2d-path vocab_mapping/t2d.npy \
  --num-workers 1 \
  --prefetch-factor 2 \
  --checkpoint-freq 2 \
  --save-best

• manual verification of checkpoint directory structure and checkpoint_best/ symlink target
• confirmed backward compatibility by running without --checkpoint-freq and --save-best (current behavior maintained)

Style/Lint Notes

• On running make style the program suggests to make modification at src/speculators/train/trainer.py:204:13. Since this is unrelated to my modifications, I have not updated it to satisfy make style.
• make style also suggests to make import blocks sorted in src/speculators/train/checkpointer.py:1:1. However, since the suggestion is also present for the current version of the repo, I have not made any changes to the arrangement of the import block.

[rahul-tuli]

I think some IDE files were also added as a part of this diff, could you remove them?

[surojitiitg]

I think some IDE files were also added as a part of this diff, could you remove them?

@rahul-tuli thanks for catching that. I have removed them. Let me know if any further modifications are required.

[rahul-tuli]

LGTM, pending comments, great first pass!

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @surojitiitg.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[fynnsu]

@surojitiitg Please fix quality issues.
You can run make style to try to auto-fix them
and make quality to confirm that they are all passing.

[fynnsu]

Looks good! Thank you for making all the updates.