Skip to content

DPO 不支持 per-pair margin / DPOTrainer does not support per-pair margin in loss #10402

Open

Description

@hejinhua1

Checklist / 检查清单

  • I have searched existing issues, and this is a new feature request. / 我已经搜索过现有的 issues,确认这是一个新的 Feature Request。

Feature Request Description / Feature Request 描述

Background / 背景

The margin column is already a first-class citizen in the Swift data
pipeline:

  • template/template_inputs.py:81 — StdTemplateInputs.margin: Optional[float]
  • template/base.py:550-565 — _rlhf_encode() writes margin into the
    encoded sample (used by both DPO and RM encoding paths)
  • template/base.py:1877-1880 — _rlhf_data_collator automatically
    collates margin into a res['margin'] tensor when present in the batch

RewardTrainer.compute_loss (reward_trainer.py:43-53) already consumes
this margin:

loss = -F.logsigmoid(rewards_chosen - rewards_rejected - margin).mean()

However, DPOTrainer.dpo_loss (dpo_trainer.py:228-363) never reads the margin column. Worse, since concatenated_forward does model(**batch, use_cache=False) without popping margin out of batch first, passing a dataset with a margin field to DPO currently crashes with TypeError: ... unexpected keyword argument 'margin'.

Pull Request / Pull Request 信息

Pull Request / Pull Request 信息

I have already implemented this fix locally and am ready to submit a PR
that links to this issue. See the three changes described above in
`swift/rlhf_trainers/dpo_trainer.py`. I used AI assistance for code
exploration and implementation; I have reviewed every changed line myself.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    enhancementNew feature or request

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions