fix(trainer): avoid KeyError when a checkpoint is saved at an unvalidated step

[KamitobiHaru]

Bug

When save_checkpoint_interval is not a multiple of validate_interval, a checkpoint can be saved at a step that was never validated. The best-ckpt comparison then indexes val_acc_step_or_epoch[ckpt_name], raising KeyError right after the checkpoint was written and aborting the whole run.

Reproduce

save_checkpoint_interval=5000, validate_interval=2000 → crash at step 5000: KeyError: 'model.pt.ep0.5000'.

Fix

Guard both the DeepSpeed and the torch.save checkpoint paths with .get(). An unvalidated checkpoint is skipped for the "best" bookkeeping (and logged) instead of crashing; saved_ckpts falls back to 0.0.

Verified

Two-stage training run (TTS 6.3k → real 517, 50 epochs) completed without the crash across ~50 checkpoint saves.

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

[LauraGPT]

The KeyError is avoided, but this fallback makes an unvalidated checkpoint participate in best-model ranking with a fabricated score of 0.0. That can delete a real validated checkpoint.

I reproduced this on commit f3c5bd73e63073063e4ca2ccf8665b4be22da32e with loss-based ranking and keep_nbest_models=1:

1. save validated model.pt.ep1.1 with loss 1.0;
2. save model.pt.ep1.2 at a step that has no validation metric;
3. the new checkpoint is inserted into saved_ckpts with 0.0;
4. the pruning code removes model.pt.ep1.1, leaving only the unvalidated checkpoint.

Please keep checkpoints without the configured validation metric out of saved_ckpts (or retain them through a separate, explicitly non-ranked policy) instead of assigning a synthetic score. Please also add focused regression coverage for both loss and acc ranking, including the DeepSpeed and non-DeepSpeed save paths, so an unvalidated checkpoint cannot evict a validated best checkpoint.

[KamitobiHaru]

@LauraGPT thanks for the review — both points are addressed.

1. Unvalidated checkpoints no longer enter ranking.

I removed the .get(ckpt_name, 0.0) fallback entirely. A checkpoint saved at a step with no configured validation metric (acc/loss) is now kept on disk but excluded from saved_ckpts, so it never competes in best-model ranking and never participates in keep_nbest_models pruning. best_step_or_epoch is also only ever set to a checkpoint that carries the metric, so it can never point at an unvalidated save.

Your repro is now the exact regression test: with keep_nbest_models=1 and loss ranking, saving validated model.pt.ep1.1 (loss 1.0) followed by unvalidated model.pt.ep1.2 leaves saved_ckpts == {"model.pt.ep1.1": 1.0} — the validated best survives and the unvalidated file stays on disk.

2. Regression coverage added in tests/test_trainer_ds_unvalidated_ckpt.py, parametrized over loss/acc × the three checkpointing paths (trainer_ds torch.save, trainer_ds DeepSpeed, and the non-DeepSpeed trainer):

• test_unvalidated_checkpoint_cannot_evict_validated_best — the scenario above, all 6 combinations.
• test_validated_checkpoints_still_rank_and_prune — control: validated checkpoints still rank and keep_nbest_models pruning still deletes the worse validated checkpoint, so normal behavior is unchanged.

While covering the non-DeepSpeed trainer, I found it shared the identical bug (direct indexing of the metric dicts), so it is fixed the same way in the same PR.

Honesty of the tests: on upstream main (before this PR), all 6 unvalidated-step cases fail with the original KeyError: 'model.pt.ep1.2'; with the fix all 12 cases pass. The DeepSpeed and torch.save paths each use on-disk checkpoint artifacts so the eviction/deletion is observable in both.

[LauraGPT]

Re-reviewed the updated head 99af35d. The requested ranking issue is fixed across all three save paths: checkpoints without the configured validation metric remain on disk but no longer enter saved_ckpts, best-model selection, or keep_nbest_models pruning.\n\nVerification on the exact head:\n- New focused matrix: 12 passed (loss/acc x Trainer, TrainerDS torch.save, TrainerDS DeepSpeed fake).\n- The same test file against the pre-fix f3c5bd7 produced 6 failures and 6 control passes, confirming the regression coverage is meaningful.\n- Broader collectable suite: 432 passed, 1 skipped, 1 unrelated environment failure in test_auto_model because this host lacks modelscope/rapidfuzz and other optional runtime dependencies.\n- compileall passed for both trainer modules; git diff --check passed.\n\nThanks for covering the non-DeepSpeed path and strengthening the DeepSpeed+acc assertion.