Packing in DPOTrainer
- 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
Assessment
This issue has not been assessed yet.