add resume for dflash training script

[tianhaoz95]

Motivation

The current dflash saves checkpoint but does not have logic to resume training, dflash tend to have a higher epoch count compared with eagle3, so this will be helpful.

Modifications

• change the save checkpoint to be on the epoch resolution since we don't have logic to restore optimizer state
• when loading drafter model, use checkpoint if requested with --resume and available

Related Issues

NA

Accuracy Test

The loading and training proceeds:

INFO:specforge.utils:/tmp/tianhaoz/longcat_dflash_training/outputs
INFO:specforge.utils:Last checkpoint detected: /tmp/tianhaoz/longcat_dflash_training/outputs/epoch_4
INFO:specforge.utils:Loaded draft config from checkpoint: /tmp/tianhaoz/longcat_dflash_training/outputs/epoch_4

Benchmark & Profiling

This change is unrelated to performance.

Checklist

• Format your code according to the Code Formatting with Pre-Commit.
• Add unit tests as outlined in the Running Unit Tests.
• Update documentation / docstrings / example tutorials as needed, according to Writing Documentation.
• Provide throughput / latency benchmark results and accuracy evaluation results as needed, according to Benchmark and Profiling and Accuracy Results.
• For reviewers: If you haven't made any contributions to this PR and are only assisting with merging the main branch, please remove yourself as a co-author when merging the PR.
• Please feel free to join our Slack channel at https://sgl-fru7574.slack.com/archives/C09784E3EN6 to discuss your PR.

[gemini-code-assist]

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!