Fix All chunked_loss Benchmark Scripts
#438
Conversation
| @@ -1,4 +1,5 @@ | |||
| from test.chunked_loss.test_dpo_loss import HF_DPO_Loss | |||
There was a problem hiding this comment.
I believe this fn wouldn't work because right now return is a tuple
https://github.com/linkedin/Liger-Kernel/blob/aeb8ce38b637eeecd92058ad6bbc0fac88fb257f/benchmark/scripts/benchmark_dpo_loss.py#L120C1-L122C21
There was a problem hiding this comment.
also I believe other chunked loss benchmarking scripts would face the same issue
There was a problem hiding this comment.
Got it! Thanks for pointing that out~
austin362667
changed the title
chunked_loss Benchmark Scripts
| @@ -65,13 +43,19 @@ def forward(self, x, target): | |||
| self.lin.weight, | |||
There was a problem hiding this comment.
Can we rather just output the loss from Torch[CUSTOM]Loss and Liger[CUSTOM]Loss, ie. just output tuple's first value. As we only need the loss, this would simplify the other code
There was a problem hiding this comment.
Sure. Good idea. @shivam15s Thanks for review!
I'm using lambda to create a function that will access the first element only when it's actually called.
Reason is that I want to reuse TorchLMHead[CUSTOM] in tests without modifying them as well as for TorchLMHead[CUSTOM].
This comment was marked as outdated.
This comment was marked as outdated.
Sorry, something went wrong.
Signed-off-by: Austin Liu <[email protected]> Fix all `.backward()` and reuse `TorchLMHead` from tests Signed-off-by: Austin Liu <[email protected]>
Signed-off-by: Austin Liu <[email protected]> Reuse LigerLMHead Signed-off-by: Austin Liu <[email protected]>
austin362667
force-pushed
the
austin362667/fix/dpo_bench
branch
from
7776009 to
291f933
Compare
Summary
With recent updates, we need to align the APIs used in the benchmark script to ensure consistency.
Code changes:
.backward()pass inbackwardandfullmode.TorchLMHeadfromchunked_losstests.Testing Done
make testto ensure correctnessmake checkstyleto ensure code stylemake test-convergenceto ensure convergence