From f01ffb1f5d3af9df407fc4f1b442f8b3dca15849 Mon Sep 17 00:00:00 2001 From: Aryan Dhawan Date: Mon, 28 Sep 2026 06:20:53 +0000 Subject: [PATCH] docs: fix input/target shapes in DPO and ORPO forward docstrings (#1406) ## Summary The `forward` docstrings of `LigerFusedLinearDPOFunction` and `LigerFusedLinearORPOFunction` describe `_input` (and DPO's `ref_input`) as `(batch_size * seq_len, hidden_size)` and `target` as `(batch_size * seq_len,)`. The implementation requires `(batch_size, seq_len, hidden_size)` and `(batch_size, seq_len)`: - `LigerFusedLinearPreferenceBase` chunks dim 0 by sequence (`torch.chunk(_input[:len_chosen], dim=0)`), reduces per-token log-probs over dim 1, and documents these shapes itself (`_compute_loss` describes `input_chunk` as `(2 * chunk_size, sequence_length, hidden_size)`). - `test/chunked_loss/test_dpo_loss.py` and `test/chunked_loss/test_orpo_loss.py` construct 3D inputs (`torch.randn(B, T, H, ...)`). - Following the docstrings as written would silently split tokens instead of sequences into chosen/rejected halves. Docstring-only change; `make checkstyle` passes; no kernel or convergence impact. The same wording also appears in `cpo_loss.py` (lines 62, 64), `simpo_loss.py` (lines 70, 72), and `kto_loss.py` (lines 94, 96, 99, including `ref_input`). Happy to fold those into this PR if you prefer one sweep. ## Testing Done - Hardware Type: n/a (docstring-only) - [x] run `make checkstyle` to ensure code style - [ ] run `make test` to ensure correctness (unaffected) - [ ] run `make test-convergence` to ensure convergence (unaffected) Co-authored-by: Vaibhav Jindal --- src/liger_kernel/chunked_loss/dpo_loss.py | 6 +++--- src/liger_kernel/chunked_loss/orpo_loss.py | 4 ++-- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/src/liger_kernel/chunked_loss/dpo_loss.py b/src/liger_kernel/chunked_loss/dpo_loss.py index 7db8c2e9..ac861ca5 100644 --- a/src/liger_kernel/chunked_loss/dpo_loss.py +++ b/src/liger_kernel/chunked_loss/dpo_loss.py @@ -178,11 +178,11 @@ class LigerFusedLinearDPOFunction(LigerFusedLinearPreferenceBase): """ Fused linear layer with DPO loss. Args: - _input (torch.Tensor): Input tensor. Shape: (batch_size * seq_len, hidden_size) + _input (torch.Tensor): Input tensor. Shape: (batch_size, seq_len, hidden_size) weight (torch.Tensor): Weight tensor. Shape: (vocab_size, hidden_size) - target (torch.LongTensor): Target tensor. Shape: (batch_size * seq_len,) + target (torch.LongTensor): Target tensor. Shape: (batch_size, seq_len) bias (torch.Tensor, optional): Bias tensor. Shape: (vocab_size,) - ref_input (torch.Tensor, optional): Reference model input tensor. Shape: (batch_size * seq_len, hidden_size) + ref_input (torch.Tensor, optional): Reference model input tensor. Shape: (batch_size, seq_len, hidden_size) ref_weight (torch.Tensor, optional): Reference model weight tensor. Shape: (vocab_size, hidden_size) ref_bias (torch.Tensor, optional): Reference model bias tensor. Shape: (vocab_size,) ignore_index (int): Index to ignore in loss computation diff --git a/src/liger_kernel/chunked_loss/orpo_loss.py b/src/liger_kernel/chunked_loss/orpo_loss.py index 7cb0bc37..d3237e1b 100644 --- a/src/liger_kernel/chunked_loss/orpo_loss.py +++ b/src/liger_kernel/chunked_loss/orpo_loss.py @@ -60,9 +60,9 @@ class LigerFusedLinearORPOFunction(LigerFusedLinearPreferenceBase): """ Fused linear layer with ORPO loss. Args: - _input (torch.Tensor): Input tensor. Shape: (batch_size * seq_len, hidden_size) + _input (torch.Tensor): Input tensor. Shape: (batch_size, seq_len, hidden_size) weight (torch.Tensor): Weight tensor. Shape: (vocab_size, hidden_size) - target (torch.LongTensor): Target tensor. Shape: (batch_size * seq_len,) + target (torch.LongTensor): Target tensor. Shape: (batch_size, seq_len) bias (torch.Tensor, optional): Bias tensor. Shape: (vocab_size,) ignore_index (int): Index to ignore in loss computation beta (float): Weight for the odds ratio loss