TokensPerSecondCallback captures resume_from_checkpoint before auto-resume detection resolves it, silently zeroing token metrics

Expected Behavior

When include_tkps: true and auto_resume_from_checkpoints: true are both set (no explicit resume_from_checkpoint in the config), resuming a run should restore the saved tokens/total and tokens/trainable counters from the checkpoint being auto-resumed, the same way it does when resume_from_checkpoint is set explicitly in the YAML.

Current behaviour

On an auto-resumed run, TokensPerSecondCallback.on_train_begin() never restores the saved token counters, so tokens/total, tokens/trainable and tokens/train_per_sec_per_gpu reset to zero/spurious values, silently, with no warning.

Root cause: get_callbacks() in src/axolotl/core/builders/causal.py constructs TokensPerSecondCallback(resume_from_checkpoint=self.cfg.resume_from_checkpoint) inside setup_model_and_trainer(), which train.py calls before determine_last_checkpoint(cfg) resolves the auto-resume path. At construction time cfg.resume_from_checkpoint is still None, and the callback stores that value directly (self.resume_from_checkpoint = resume_from_checkpoint), so on_train_begin()'s isinstance(self.resume_from_checkpoint, str) check fails and the restore is skipped, even though cfg.resume_from_checkpoint gets the correct path a few lines later in train.py. Passing resume_from_checkpoint explicitly in the config works fine, since it's already a string at construction time.

I confirmed this by calling the real determine_last_checkpoint() and TokensPerSecondCallback directly, in the same order train.py calls them: with a checkpoint directory containing a tokens_state.json, constructing the callback before calling determine_last_checkpoint(cfg) leaves state.tokens at None after on_train_begin(); constructing it after (with the resolved path) restores state.tokens correctly. I did not run a full GPU training job through an actual resume cycle, only this direct call sequence against main.

Steps to reproduce

1. Set include_tkps: true and auto_resume_from_checkpoints: true, with no resume_from_checkpoint in the config.
2. Train until a checkpoint is saved (which also writes tokens_state.json).
3. Restart the same run so it auto-resumes from that checkpoint.
4. Compare the first few tokens/total / tokens/trainable log values against the checkpoint's saved tokens_state.json - they don't carry over.

Config yaml

include_tkps: true
auto_resume_from_checkpoints: true
# resume_from_checkpoint intentionally omitted, relying on auto-resume

Possible solution

Resolve determine_last_checkpoint(cfg) before setup_model_and_trainer() runs, or have get_callbacks() re-read self.cfg.resume_from_checkpoint at on_train_begin() time instead of at construction time. Happy to send a PR for either, whichever fits the codebase better.

Which Operating Systems are you using?

• Linux

Python Version

3.11 (bug is in the call ordering, not Python-version-dependent)

axolotl branch-commit

main/e198b1d8ed1d36b8d732f67cbb1dda62f86d3986

Acknowledgements

• My issue title is concise, descriptive, and in title casing.
• I have searched the existing issues to make sure this bug has not been reported yet.
• I am using the latest version of axolotl.
• I have provided enough information for the maintainers to reproduce and diagnose the issue.