Conversation
Adds three new reward types to the online-data-mixing plugin's compute_reward dispatch, ported from a learnability/velocity-driven data mixing research prototype: - LEARNABILITY: 1 - loss_few_shot/loss_zero_shot, comparing a category's zero-shot vs few-shot (templated) eval loss. - VELOCITY: 1 - current_loss/previous_loss, tracking how quickly a category's eval loss is still dropping (module-level buffer, no dev set needed beyond the existing eval_dataset_dict). - COMBINED: exponentially-decayed blend of the two, alpha(t) = exp(-beta*t/total_steps), favoring learnability early in training and velocity later. OnlineMixingDataset gains optional templated_eval_dataset_dict / templated_eval_collators_dict / beta constructor params (required together with LEARNABILITY/COMBINED, validated at construction time) and threads them through _reset_eval_dataloaders, update_sampling_weights, and _extract_information_from_state_for_reward. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> Signed-off-by: Amal Joe <amal.joe@ibm.com>
amaljoe
force-pushed
the
odm-learnability-velocity-rewards
branch
from
September 1, 2026 06:49
19445bc to
1931379
Compare
Splits compute_reward() and _extract_information_from_state_for_reward() into per-reward-type helper functions plus a dispatch table, bringing each back under pylint's too-many-return-statements threshold (both had grown to 9 returns after the LEARNABILITY/VELOCITY/COMBINED additions). Also fixes a not-an-iterable false positive in test_online_data_update_sampling_weights_with_templated_eval_dataset: OnlineMixingDataset.log["count"] starts as an int literal in dataset.py before being overwritten with a list by update_sampling_weights(), so pylint's static inference can't see it's iterable by the time the test reads it. Wrapping in list() satisfies the checker without changing runtime behavior. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> Signed-off-by: Amal Joe <amal.joe@ibm.com>
update_sampling_weights() falls back to torch.device(0) (cuda:0) when no accelerator is passed, so the test previously moved the model there too. That happens to resolve to mps:0 on Apple Silicon but crashes with "Found no NVIDIA driver" on the CPU-only CI runner, since torch.device(0) always means CUDA regardless of platform. Pass a minimal single-process accelerator stub (device=cpu, no-op prepare/reduce) instead, so the test exercises the same accelerator-driven code path update_sampling_weights() takes under real (single-process) Accelerate usage, without depending on any GPU/MPS backend being present. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> Signed-off-by: Amal Joe <amal.joe@ibm.com>
kmehant
self-requested a review
September 8, 2026 17:16
kmehant
approved these changes
Sep 8, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
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.
Summary
Rewardtypes to the online-data-mixing plugin'scompute_rewarddispatch:LEARNABILITY(1 - loss_few_shot/loss_zero_shot),VELOCITY(1 - current_loss/previous_loss, tracked via a module-level buffer), andCOMBINED(exponentially-decayed blend of the two,alpha(t) = exp(-beta*t/total_steps))OnlineMixingDatasetgains optionaltemplated_eval_dataset_dict/templated_eval_collators_dict/betaconstructor params — the first two are required together whenreward_typeisLEARNABILITYorCOMBINED(validated at construction, raisesValueErrorotherwise), and are threaded through_reset_eval_dataloaders,update_sampling_weights, and_extract_information_from_state_for_rewardTest plan
pytest tests/test_compute_reward.py -v— 9/9 passed (3 new: learnability, velocity, combined)pytest tests/test_online_data.py -v— 4/4 passed (2 new: validation error, end-to-endupdate_sampling_weights)pytest tests/— 19/19 passed, no regressions