huggingface / huggingface/trl

Packing in DPOTrainer

Open
#2,469 5 comments 0 reactions 0 assignees View on GitHub
✨ enhancement 🏋 DPO 🚨 Important
Dominant language
Python
Stars
19.3k
Forks
3k
Avg merge
1d 20h
Merged PRs (30d)
194

Description

### Feature request

packing can be supported in dpo trainer.

### Motivation

sequences with high variation in length padded together will waste a lot of resources. issue #1274 mentioned this, but i don't think the conclusion is correct. packing won't conflict with pairwise datasets. you just need to unpack the sequence after forwarding. you can easily identify the start and the end of each sequence by the position of 0 in `position_ids`.

### Your contribution

actually i already have a version of dpo trainer that can deal with packing:

```python
def _forward_one(self, model: nn.Module, batch: dict, name: str):
batch = dict(batch)
targets = batch.pop("targets")
loss_masks = batch.pop("loss_masks").to(torch.bool)
out = model(**batch)
logits = out["logits"]
targets = targets.clone()
targets[~loss_masks] = 0 # dummy token; we'll ignore the losses on these tokens later
per_token_logps = torch.gather(logits.log_softmax(-1), dim=2, index=targets.unsqueeze(2)).squeeze(2)
per_token_logps[~loss_masks] = 0
# unpack logp
assert batch["position_ids"].shape[0] == 1
starts = (batch["position_ids"] == 0).nonzero()[:, 1]
seqs = []
for i in range(len(starts) - 1):
seqs.append(per_token_logps[0, starts[i]: starts[i + 1]].sum())
seqs.append(per_token_logps[0, starts[-1]:].sum())
all_logps = torch.stack(seqs)
return {
f"{name}_logps": all_logps,
f"mean_{name}_logits": logits[loss_masks].mean(),
}

def concatenated_forward(self, model: nn.Module, batch: dict):
ret = {}
for name in ("chosen", "rejected"):
ret = {**ret, **self._forward_one(model, batch[name], name)}
return ret
```

this does not directly fit into the complex logic of dpo trainer, but the idea is that it is possible to do packing.

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.