[Swanand-Kadhe]

This PR ports changes from finetune to enable DPO

• DPO data collator is copied from padding-free branch
• Running this branch is giving the following error:

Reproduces previous results before moving to main

• matches the previous result using dpo
• ran a new result using dpo_norm, looks like it doesnt perform as well as the latter (perhaps it requires new settings for beta)

[garrett361]

Didn't we merge fix/dataset into main in #27 ? Should this PR be against main @Swanand-Kadhe ?

[garrett361]

The first error is something subtle with the @dataclass decorator around the TensorDataCollatorWithFlatteningDPO subclass. Removing the decorator removes the error for me.

There's not actually any need for the TensorDataCollatorWithFlattening superclass to inherit from DefaultDataCollator or to be a dataclass. So I'd suggest removing both of these, and it should make everything easier.

[garrett361]

EDIT: the below is wrong, I misread the code. filter_batch only acts on keys, not values.

This TensorDataCollatorWithFlatteningDPO class doesn't seem like it will work, though @Swanand-Kadhe , because by the time the collator is used, the data has already been tokenized. So, the filter_batch function only acts on tensors, not on strings.

[garrett361]

Added a test in #29 , though I'm not sure it tests the right things. Please let me know @Swanand-Kadhe

[fabianlim]

@garrett361 I made some fixes and i tested it and it should be working now.

• After dataset transformation, the keys will be dict_keys(['chosen_input_ids', 'chosen_labels', 'chosen_attention_mask', 'rejected_input_ids', 'rejected_labels', 'rejected_attention_mask'])
• DataLoader will load the data, and apply the TensorDataCollatorWithFlatteningDPO, which basically applys TensorDataCollatorWithFlattening to both chosen_ and rejected_ groups. The cu_seq_len's etc will be generated for both groups.
• then in dpo_utils.concatenated_forward we will dispatch based on padding_free=True to padding_free.concatenated_inputs, which we then concat the input_ids, labels, cu_seq_len, etc

[garrett361]

Sounds good, do you mind adding an actual test @fabianlim ? Can build off of #29?

[garrett361]

just noting that we'll want this to be padding_free --> packing when the pr is made against main

[fabianlim]

this is actually an internal flag. the one in FlatArguments has already been changed.. but ok I changed

[garrett361]

yeah, just wanted to make it less confusing for future developers anyway.

[fabianlim]

@garrett361 @Swanand-Kadhe it matches my previous results, I updated the main comment

[Swanand-Kadhe]

Thank you very much, @fabianlim and @garrett361!
Shall we merge the PR so that changes are in the branch? I can then create a PR to merge these changes to main. If you have other suggestions, let me know.

[fabianlim]

@Swanand-Kadhe we already changed this PR to merge with main and it was already rebased

[Swanand-Kadhe]

Ah, right, sorry for missing it @fabianlim. It would be good to merge it with main. In that way, we can use main for DPO + SFT. Many thanks.

[fabianlim]

We have updated this PR and fixed the issues. merging