DPO and KTO skip DDP gradient synchronization with use_liger_kernel=True

unwrapped_model = self.accelerator.unwrap_model(model)
if is_zero3 or self.is_fsdp_enabled:
    return self._forward_redirection(model, unwrapped_model, self._compute_loss, unwrapped_model, inputs, return_outputs)
return self._compute_loss(unwrapped_model, inputs, return_outputs)

# ddp_grad_sync.py
import os

import torch
import torch.distributed as dist
from datasets import Dataset
from transformers import AutoModelForCausalLM, AutoTokenizer, TrainerCallback

from trl import DPOConfig, DPOTrainer

MODEL = "trl-internal-testing/tiny-Qwen3ForCausalLM"
USE_LIGER = os.environ.get("USE_LIGER", "1") == "1"
rank = int(os.environ["RANK"])

tokenizer = AutoTokenizer.from_pretrained(MODEL)
# Each rank gets different data, so unsynchronised gradients are guaranteed to differ.
dataset = Dataset.from_list(
    [
        {
            "prompt": f"rank {rank} question {i} about something",
            "chosen": f"a good answer from rank {rank}",
            "rejected": f"a bad answer from rank {rank}",
        }
        for i in range(32)
    ]
)
model = AutoModelForCausalLM.from_pretrained(MODEL, dtype=torch.bfloat16)
training_args = DPOConfig(
    output_dir="ddp_grad_sync_out",
    use_liger_kernel=USE_LIGER,
    per_device_train_batch_size=2,
    max_steps=1,
    max_length=128,
    bf16=True,
    save_strategy="no",
    report_to="none",
)
trainer = DPOTrainer(model=model, args=training_args, train_dataset=dataset, processing_class=tokenizer)


class GradProbe(TrainerCallback):
    def on_pre_optimizer_step(self, args, state, control, **kwargs):
        unwrapped = trainer.accelerator.unwrap_model(trainer.model)
        n_differing, n_total, max_diff = 0, 0, 0.0
        for _, param in unwrapped.named_parameters():
            if param.grad is None:
                continue
            grad = param.grad.detach().float()
            gathered = [torch.zeros_like(grad) for _ in range(dist.get_world_size())]
            dist.all_gather(gathered, grad)
            diff = (gathered[0] - gathered[1]).abs().max().item()
            n_total += 1
            max_diff = max(max_diff, diff)
            if diff > 0:
                n_differing += 1
        if dist.get_rank() == 0:
            print(f"use_liger_kernel={USE_LIGER} params={n_total} differing={n_differing} max_abs_diff={max_diff:.3e}")


trainer.add_callback(GradProbe())
trainer.train()

USE_LIGER=1 torchrun --nproc_per_node=2 ddp_grad_sync.py
USE_LIGER=0 torchrun --nproc_per_node=2 ddp_grad_sync.py

use_liger_kernel=True  params=25 differing=25 max_abs_diff=1.265e-01
use_liger_kernel=False params=25 differing=0  max_abs_diff=0.000e+00