Skip to content

fix(train): replay epochs after prevalidation rollback - #273

Merged
isLinXu merged 1 commit into
Tencent:mainfrom
onion-hong:codex/fix-prevalidation-replay
Sep 11, 2026
Merged

fix(train): replay epochs after prevalidation rollback#273
isLinXu merged 1 commit into
Tencent:mainfrom
onion-hong:codex/fix-prevalidation-replay

Conversation

@onion-hong

Copy link
Copy Markdown
Contributor

When an epoch-end callback makes the online model non-finite, _recover_before_validation() restores the healthy checkpoint but returns false. Validation then succeeds on restored weights, so the rolled-back attempt is logged as a completed epoch, the recovery counter is cleared, and the MoE scheduler receives recovered=False. A repeated fault can therefore silently discard each epoch without exhausting the retry limit.

Return the replay signal after an online rollback, reset gradient accumulation before replay, and finalize that attempt with recovered=True, validated=False. EMA-only resynchronization keeps its existing behavior. Adjust the existing online/EMA test to inject once; add offline real-training regression cases for one-time and persistent faults, checking optimizer updates, validation calls, CSV rows, scheduler flags and the recovery budget.

Validation:

  • New regression cases fail on the original code (no replay; persistent faults do not raise).
  • On isolated upstream 3bd0a60 plus this patch: recovery/NaN/MoE scheduler tests: 72 passed, 1 skipped.
  • Related engine EMA-resync and online/EMA recovery tests: 2 passed.
  • Ruff format and git diff --check pass; modified trainer retains 11 pre-existing lint diagnostics.

This fixes prevalidation rollback accounting; it does not prevent the original source of non-finite values.

@isLinXu
isLinXu merged commit cb82422 into Tencent:main Sep 11, 2026
11 of 17 checks passed
@isLinXu

isLinXu commented Sep 11, 2026

Copy link
Copy Markdown
Collaborator

LGTM

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants