Introduce prepare for eval, fix evaluation bug #789
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Description
This pull request introduces a
prepare_for_eval
function and updates the code to support it.The implementation follows the blueprint of @fsschneider in #719 (comment) and fixes the bug of giving a free evaluation to a submission that goes out of max_runtime (again #719 (comment)).
Function signature
The arguments of
prepare_for_eval
are the same asupdate_params
, except forbatch
. I believe thatprepare_for_eval
should indeed be agnostic to the last batch used during training. The return type is the same asupdate_params
.List of changes
In
submission_runner.py
:prepare_for_eval
profiler
del batch
beforeprepare_for_eval
(instead than before evaluation)accumulated_submission_time
afterprepare_for_eval
is_time_remaining
afterprepare_for_eval
is_time_remaining
prep_eval_rng
Minor changes:
PrepareForEvalFn
tospec
prepare_for_eval
to submission templateprepare_for_eval
to all pytorch and jax submissionsFixes #719 and #758 .