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 <[email protected]>

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

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