Precompute DPO logprobs

[natolambert]

To save compute.
Another hard issue :)

[natolambert]

A sketch of how this could work.

Add an option in dpo_tune where instead of using concatenated_forward, we run just forward for each with an optional save of the logprobs.

open-instruct/open_instruct/dpo_tune.py

Line 568 in 42c1fa3

for epoch in range(starting_epoch, args.num_train_epochs):

Then, you iterate over batches and compute loss and update the model.

Optional: logic to move one model into cuda at a time. Shouldn't be too hard.

[hamishivi]

Yeah, sounds about right. It's very easy to implement, I did it in EasyLM (although sharding issues mean its broken), but the logic should be right: https://github.com/hamishivi/EasyLM/blob/main/EasyLM/models/llama/llama_train_dpo.py#L372-L400