Skip to content

docs: fix input/target shapes in DPO and ORPO forward docstrings - #1406

Open
Aryan95614 wants to merge 1 commit into
linkedin:mainfrom
Aryan95614:docs/dpo-orpo-input-shapes
Open

Aryan95614 wants to merge 1 commit into
linkedin:mainfrom
Aryan95614:docs/dpo-orpo-input-shapes

Conversation

@Aryan95614

Copy link
Copy Markdown

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)
  • run make checkstyle to ensure code style
  • run make test to ensure correctness (unaffected)
  • run make test-convergence to ensure convergence (unaffected)

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant