[Liger] liger DPO support

[kashif]

What does this PR do?

Add support for Liger-kernel losses for the DPO Kernel

Needs: linkedin/Liger-Kernel#521

Peft support: #3065

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[qgallouedec]

Can you use @parameterized.expand instead?

[qgallouedec]

liger loss isn't compatible with ref precomputing right? If so we could add a warning or an error.

[qgallouedec]

@kashif I've added this section in the new guide for reducing memory usage, if you've words to fill it

[VProv]

This PR needs to use _FSDPForwardRedirection or another solution to work with FSDP correctly
linkedin/Liger-Kernel#615
https://github.com/linkedin/Liger-Kernel/blob/2bb8dcfc18f10ff90f942f238b5cfe16c12749b6/src/liger_kernel/transformers/trainer/orpo_trainer.py#L18-L66

[kashif]

@VProv, at the moment, I was having issues getting the same outputs/metrics with and without liger in the trainer.

[VProv]

@VProv, at the moment, I was having issues getting the same outputs/metrics with and without liger in the trainer.

What setup are you using?

[vaibhavjindal]

Hi, I am working on fixing the output/metrics issue.
Added a PR in liger-kernel: linkedin/Liger-Kernel#676

[vaibhavjindal]

@kashif @qgallouedec can you please review the following PR which fixes the output/metrics issue? Thanks :)
#3346

[kashif]

thanks @vaibhavjindal done, i'll fix the merge conflict and then review this PR

[hanbyul-kim]

Hi, thanks for sharing your work! Can I use your code with DeepSpeed Zero 3? I tried running it with that setup, but it doesn't seem to be working. I think it's related to parameter partitioning based on my analysis of the error log.

[rank5]:   File "/mnt/nappipe/users/hanbyul-kim/RORL/apply_liger_loss/Liger-Kernel/src/liger_kernel/chunked_loss/dpo_loss.py", line 94, in forward
[rank5]:     return super().forward(
[rank5]:   File "/mnt/nappipe/users/hanbyul-kim/RORL/apply_liger_loss/Liger-Kernel/src/liger_kernel/chunked_loss/fused_linear_preference.py", line 241, in forward
[rank5]:     accumulate_chunk(input_chunk, target_chunk, ref_input_chunk, chosen_nll_target_chunk)
[rank5]:   File "/mnt/nappipe/users/hanbyul-kim/RORL/apply_liger_loss/Liger-Kernel/src/liger_kernel/chunked_loss/fused_linear_preference.py", line 159, in accumulate_chunk
[rank5]:     ) = fused_fwd_bwd(input_chunk, target_chunk, ref_input_chunk, chosen_nll_target_chunk)
[rank5]:   File "/mnt/nappipe/users/hanbyul-kim/RORL/apply_liger_loss/Liger-Kernel/src/liger_kernel/chunked_loss/fused_linear_preference.py", line 120, in fused_fwd_bwd
[rank5]:     return torch.func.grad_and_value(compute_loss, argnums=(0, 1), has_aux=True)(
[rank5]:   File "/root/.dpo_trainer_venv/lib/python3.10/site-packages/torch/_functorch/apis.py", line 440, in wrapper
[rank5]:     return eager_transforms.grad_and_value_impl(
[rank5]:   File "/root/.dpo_trainer_venv/lib/python3.10/site-packages/torch/_functorch/vmap.py", line 48, in fn
[rank5]:     return f(*args, **kwargs)
[rank5]:   File "/root/.dpo_trainer_venv/lib/python3.10/site-packages/torch/_functorch/eager_transforms.py", line 1409, in grad_and_value_impl
[rank5]:     output = func(*args, **kwargs)
[rank5]:   File "/mnt/nappipe/users/hanbyul-kim/RORL/apply_liger_loss/Liger-Kernel/src/liger_kernel/chunked_loss/fused_linear_preference.py", line 377, in _compute_loss
[rank5]:     ) = LigerFusedLinearPreferenceBase.chunk_forward(
[rank5]:   File "/mnt/nappipe/users/hanbyul-kim/RORL/apply_liger_loss/Liger-Kernel/src/liger_kernel/chunked_loss/fused_linear_preference.py", line 289, in chunk_forward
[rank5]:     logits_chunk = input_chunk @ weight.t()
[rank5]: RuntimeError: size mismatch, got input (322), mat (322x4096), vec (0)

[hanbyul-kim]

Continuing my analysis, I can confirm that it's definitely connected to DeepSpeed zero 3. When I switched to stage 2, it ran smoothly without any issues.

[kashif]

thanks @hanbyul-kim for the report