fix(finetune): all_reduce val loss across devices for initial and final validation#2266
Open
discobot wants to merge 1 commit into
Open
fix(finetune): all_reduce val loss across devices for initial and final validation#2266discobot wants to merge 1 commit into
discobot wants to merge 1 commit into
Conversation
…al validation The initial and final validation in the four finetune scripts logged the rank-local loss without reducing it across devices, unlike the periodic validation. The periodic validation also reduced into a temporary tensor that was never assigned back to val_loss, so the training progress lines kept reporting the rank-local value. Reduce the loss in all three places, using the tensor returned by Fabric.all_reduce, and reuse it in the progress lines. All reductions are no-ops on a single device. Adds a CPU regression test per script that stubs Fabric.all_reduce and asserts every printed and logged val loss is the reduced value.
discobot
requested review from
andyland,
k223kim,
lianakoleva and
t-vi
as code owners
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
1 participant
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.
Fixes #2116.
The missing reduction for the initial and final validation is not limited to
lora.pyandfull.py—adapter.pyandadapter_v2.pyhave the same three validation paths, so this fixes all four finetune scripts.The stale
val_lossin the training progress line has a concrete mechanism: the periodic eval reduces into a temporaryval_loss_tensorand never assigns it back, so theval:column keeps showing the rank-local value from the latestvalidate()call. The fix reduces via the tensor returned byFabric.all_reduce(the documented reduced result) and reassigns it toval_loss, so the printed and logged values agree everywhere. Atworld_size == 1the reduction is a no-op, so single-device runs are unchanged.Each script gets a CPU regression test that stubs
Fabric.all_reduceto make the reduced loss distinguishable from the rank-local one and asserts that the initial, periodic, and final validation losses — as well as theval:column of the progress lines — all report the reduced value. The new tests fail onmainand pass with this change; the fulltests/test_lora.py,tests/test_full.py,tests/test_adapter.py, andtests/test_adapter_v2.pysuites pass locally. I also verified thevalidate→all_reducepattern in a 2-process gloo CPU run: both ranks report the cross-rank mean.