Repository navigation
feat(chunked_loss): add fused linear IRPO (Iterative Reasoning Preference Optimization) loss - #1586
Open
Yulong-Cauli wants to merge 2 commits into
Open
Yulong-Cauli wants to merge 2 commits into
Yulong-Cauli wants to merge 2 commits into
Conversation
…ence Optimization) loss Add LigerFusedLinearIRPOLoss, a chunked fused-linear implementation of the Iterative Reasoning Preference Optimization objective (arXiv 2404.19733): a DPO preference term over summed sequence log-probabilities with a reference model (the previous iteration's model M_t in the paper), plus a length-normalized NLL term on the chosen response composed by LigerFusedLinearPreferenceBase (compute_nll_loss=True, alpha=1.0 and beta=0.1 defaults per the paper). The implementation follows LigerFusedLinearDPOFunction on the LigerFusedLinearPreferenceBase and includes: - chunked_loss/irpo_loss.py (Function + nn.Module, functional alias) - correctness tests vs an HF-style reference (test_irpo_loss.py) - benchmark script (benchmark_irpo_loss.py) - README / docs table entries Part of the FlexChunkLoss alignment roadmap (linkedin#371).
Yulong-Cauli
force-pushed
the
feat/irpo-loss
branch
from
November 7, 2026 05:23
ee9f951 to
dd810ff
Compare
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Adds
LigerFusedLinearIRPOLoss, a chunked fused-linear implementation of the Iterative Reasoning Preference Optimization objective (arXiv 2404.19733), completing the IRPO item of the FlexChunkLoss alignment roadmap (#371).Formula (Equation 1 of the paper)
L_IRPO = L_DPO + α · L_NLL, with defaults β=0.1 and α=1.0 per the paper:−log σ(β·(log[πθ(y_w)/πref(y_w)] − log[πθ(y_l)/πref(y_l)]))over summed (length-unaveraged) sequence log-probabilities, with the reference model participating (the previous iteration's model M_t in the paper)LigerFusedLinearPreferenceBase(compute_nll_loss=True)The iterative parts (rejection sampling between iterations, conditional-SFT prefixing) live in the training loop; this module computes the per-batch loss given a policy and an optional reference.
Implementation
Follows
LigerFusedLinearDPOFunctiononLigerFusedLinearPreferenceBase:chunked_loss/irpo_loss.py:LigerFusedLinearIRPOFunction+LigerFusedLinearIRPOLossliger_fused_linear_irpochunked_loss/__init__.py, documented in the README / docs alignment-kernel tablestest/chunked_loss/test_irpo_loss.py(referenceHFIRPOLoss, large/small shapes, bf16/fp32, bias/ref_bias, ignore_index/β variants, functional contract, frozen weights, α-scales-NLL check)benchmark/scripts/benchmark_irpo_loss.pyVerification
ruff check+ruff format(pinned v0.14.11) cleanaverage_log_prob=True, consistent with the other paired-loss test wrappers. With summed log probabilities at large shapes in bf16, the HF-reference comparison is seed-unstable (chunked-GEMM rounding amplified by ~T; verified that existing DPO behaves identically in the previously untested{NLL + ref + summed-logp}combination). The paper-faithful summed path (the module default) is covered by the functional tests in both dtypes.Refs #371