Fix TrackioCallback fails to log evaluation metrics after training ends#46935
Conversation
lewtun
changed the title
TrackioCallback fails to log evaluation metrics after training ends
|
The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update. |
| is_torch_available, | ||
| ) | ||
| from transformers.integrations.integration_utils import KubeflowCallback, SwanLabCallback | ||
| from transformers.integrations.integration_utils import KubeflowCallback, SwanLabCallback, TrackioCallback |
There was a problem hiding this comment.
I added a regression test since we didn't have any at all for trackio. Happy to remove it if that's preferred
|
The failing tests seem unrelated to my changes |
qgallouedec
left a comment
qgallouedec
left a comment
There was a problem hiding this comment.
Didn't try myself but the justification and fix lgtm
| callback.on_train_end(args, state, control, model=model) | ||
| callback.on_log(args, state, control, model=model, logs={"eval_loss": 0.5}) | ||
|
|
||
| self.assertEqual(fake_trackio.init.call_count, 2) |
There was a problem hiding this comment.
Ah because it's called once on .setup() and then again on .on_log() because .on_train_end() sets self._initialized = False
Happy to change to test to something more explicit / instructive - I took a best guess here :)
There was a problem hiding this comment.
Not a big deal at all but maybe we just move the .on_log() above the .on_train_end(), which would be the more standard usage anyways?
There was a problem hiding this comment.
Oh I was trying to simulate the scenario where a script has finished training and then afterwards tries to log a new value to the same trackio job (e.g. running trainer.evaluate() after trainer.train()). That's why the .on_log() appears after .on_train_end() but let me add an explicit log step for training to make it clearer
abidlabs
left a comment
abidlabs
left a comment
There was a problem hiding this comment.
Looks very reasonable! Just a qq above
Assert both Trackio log calls made around train end reinitialization. Co-authored-by: OpenAI Codex <[email protected]>
CI recapDashboard: View test results in Grafana |
* upstream/main: (39 commits) Remove deprecated training args and `is_fast` property (huggingface#46917) Consistent output shape from `get_image_features` (huggingface#46405) Fix multi-device mxfp4 dequantization race in `_convert_moe_packed_tensors` (huggingface#47423) fix failed test cases for qwen3_omni_moe model (huggingface#47449) Fix Hunyuan-VL PIL image resize parity with reference preprocessing (huggingface#47233) Move `value` padding into the attention interfaces that need it (huggingface#47451) Simplify function dispatch for linear attention (huggingface#47450) [cache] Allow sliding window layers to be roll-backed for speculative decoding (huggingface#47447) Fix double-shifted training loss in GitForCausalLM (huggingface#47395) Fix CohereASR training-loss double-shift (same as Moonshine fix huggingface#46784) (huggingface#46895) Warn when `group_by_length` is silently ignored for iterable datasets (huggingface#47379) Update bug report list (huggingface#46607) Fix shape mismatch in KyutaiSpeechToText `generate()` last window (huggingface#46952) Optimize flash attention max seqlen computation in vision attention (huggingface#47170) fix: remove unreachable return in special token builder (huggingface#47420) Add Harry to slow CI (huggingface#47454) BLT: vectorize patch length processing (huggingface#47385) Fix `TrackioCallback` fails to log evaluation metrics after training ends (huggingface#46935) [Kimi] add integration tests (huggingface#47383) Fix typo in `MusicgenForCausalLM.generate()` (huggingface#46974) ...
5 participants
What does this PR do?
This PR fixes the following issue:
TrackioCallback.on_train_end()callstrackio.finish(), which clears Trackio’s active run. However, the callback keeps_initialized=True. If user code callstrainer.evaluate()ortrainer.predict()aftertrainer.train(), the callback later receiveson_log()/on_predict(). Because_initializedis still True, it skipssetup()and callstrackio.log()without an active Trackio run, raising:The fix is to reset
TrackioCallback._initializedaftertrackio.finish()inon_train_end(). Then subsequent logging re-runssetup()and reinitializes/resumes the Trackio run before logging metrics.I've also added a regression test which triggers the problem on
mainand verified the fix resolves it.Code Agent Policy
The Transformers repo is currently being overwhelmed by a large number of PRs and issue comments written by
code agents. We are currently bottlenecked by our ability to review and respond to them. As a result,
we ask that new users do not submit pure code agent PRs at this time.
You may use code agents in drafting or to help you diagnose issues. We'd also ask autonomous "OpenClaw"-like agents
not to open any PRs or issues for the moment.
PRs that appear to be fully agent-written will probably be closed without review, and we may block users who do this
repeatedly or maliciously.
This is a rapidly-evolving situation that's causing significant shockwaves in the open-source community. As a result,
this policy is likely to be updated regularly in the near future. For more information, please read
CONTRIBUTING.md.Before submitting
Pull Request checks?
to it if that's the case.
Who can review?
Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.