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!
| 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
| # TVD(P || Q) = 0.5 * |P - Q| | ||
| tv_loss = 0.5 * tl.abs(p - q) | ||
|
|
||
| grad_res = tl.where(p > q, 0.5, -0.5) |
There was a problem hiding this comment.
since we're doing gradients calculation in forward pass already, we can divide gradients by BT (BT * V) based on the reduction mode here to avoid extra calculations in backward pass and saving reduction mode in ctx
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