Add TVD Loss Kernel #324
Add TVD Loss Kernel #324
Conversation
|
@ByronHsu @qingquansong @lancerts Please let me know if any changes are required. |
| pytest.param( | ||
| torch.bfloat16, | ||
| 1e-8, | ||
| 5e-2, |
There was a problem hiding this comment.
could you help to experiment what is the lowest rtol that will not fail this test for bf16? Thanks!
There was a problem hiding this comment.
@yundai424 5e-2 is the lowest that does not fail, and with a lower step, approx 5.2e-18 is the lowest rtol that does not fail.
Test passed with rtol=5.0e-02
Testing with rtol=4.5e-02...
Test passed with rtol=4.5e-02
Testing with rtol=4.0e-02...
Test passed with rtol=4.0e-02
Testing with rtol=3.5e-02...
Test passed with rtol=3.5e-02
Testing with rtol=3.0e-02...
Test passed with rtol=3.0e-02
Testing with rtol=2.5e-02...
Test passed with rtol=2.5e-02
Testing with rtol=2.0e-02...
Test passed with rtol=2.0e-02
Testing with rtol=1.5e-02...
Test passed with rtol=1.5e-02
Testing with rtol=1.0e-02...
Test passed with rtol=1.0e-02
Testing with rtol=5.0e-03...
Test passed with rtol=5.0e-03
Testing with rtol=5.2e-18...
Test passed with rtol=5.2e-18
Testing with rtol=-5.0e-03...
FAILED
| from liger_kernel.transformers.tvd import LigerTVDLoss | ||
|
|
||
|
|
||
| class TorchTVDLoss(torch.nn.Module): |
There was a problem hiding this comment.
I feel it'll be very helpful if we can add ignore index along with this PR to make TVD complete, similar to how JSD is doing it -- https://github.com/linkedin/Liger-Kernel/blob/main/src/liger_kernel/ops/jsd.py
There was a problem hiding this comment.
+1 which would be very helpful to cover broader use cases
There was a problem hiding this comment.
Thanks for the efforts! Could you also add this to the init function in transformers folder as well same as JSD? https://github.com/linkedin/Liger-Kernel/blob/main/src/liger_kernel/transformers/__init__.py#L10
Summary
Resolves #281. Implements the TVD (Total Variation Distance) kernel by computing both the loss and gradient in the forward pass.
Testing Done
Implemented tests to verify that the results of the forward and backward passes match the Torch implementation. Additionally, added a script to benchmark the memory usage and speed of the Liger implementation compared to Torch, with the results shown below.
make testto ensure correctnessmake checkstyleto ensure code stylemake test-convergenceto ensure convergence