kwongfuk/codex-skills/exact-training-resume-guard

Exact Training Resume Guard

Use this skill whenever a task writes or edits training launchers, training entrypoints, or checkpoint logic.

This is not "load model weights and start over." The requirement is exact resume:

• continue from the last saved step
• preserve optimizer state
• preserve scheduler position
• preserve RNG state
• preserve framework and distributed state such as DeepSpeed sharded checkpoints

Hard Requirements

1. Save full training state, not model-only checkpoints.

• Do not enable options like save_only_model=true for long-running training unless the user explicitly accepts losing exact resume.
• Save model, optimizer, scheduler, scaler if used, RNG state, and distributed/framework state needed by the runtime.

1. Resume from the latest complete checkpoint automatically.

• On startup, detect the newest complete checkpoint under the resume directory.
• Resume only from checkpoints that contain every file required by the framework.
• If a checkpoint is incomplete, skip it and log that decision.

1. Handle interruption safely.

• Trap SIGTERM and SIGINT.
• Forward termination to the trainer so it can finish or flush the current save path.
• Emit the latest resumable checkpoint path during shutdown.

1. Keep checkpoints on durable storage.

• Checkpoint trees and any checkpoint-writing output_dir must be on durable storage such as /scratch, not ephemeral local scratch and not mixed into lightweight repo logs.

Implementation Checklist

• Launcher defaults:
• full-state checkpoint saving enabled
• periodic save interval defined
• resume path configurable but defaulting to the task checkpoint directory
• Entry script:
• accepts resume_from_checkpoint
• passes it into the trainer/runtime
• does not silently ignore resume
• Checkpoint validation:
• verify the framework-specific state files exist before resuming
• examples: optimizer, scheduler, RNG, trainer state, DeepSpeed shard state
• Logging:
• print startup resume source
• print skipped incomplete checkpoints
• print latest resumable checkpoint on shutdown

Framework Guidance

• Hugging Face Trainer / TRL:
• prefer trainer.train(resume_from_checkpoint=...)
• keep full checkpoint directories; do not rely on save_model() alone
• DeepSpeed:
• require optimizer/model shard state under the checkpoint step directory
• do not call a checkpoint "resumable" if only safetensors weights exist
• Accelerate or custom loops:
• persist optimizer, scheduler, scaler, RNG, and dataloader/step counters when applicable

Review Rule

A training script is not done until exact resume has been checked explicitly. If the script only reloads model weights, describe it as restart from weights, not resume training.