Repository: hiyouga/LLaMA-Factory
File: src/llamafactory/train/kto/trainer.py (Python)

Code before:

        self.finetuning_args = finetuning_args
        self.reference_free = False
        self.use_dpo_data_collator = True  # hack to avoid warning
        self.generate_during_eval = False  # disable at evaluation
        self.label_pad_token_id = IGNORE_INDEX
        self.padding_value = 0
        self.is_encoder_decoder = model.config.is_encoder_decoder
        self.precompute_ref_log_probs = False
        self._precomputed_train_ref_log_probs = False
        self._precomputed_eval_ref_log_probs = False
        self._peft_has_been_casted_to_bf16 = False

        self.ref_model = ref_model
        self._stored_metrics = defaultdict(lambda: defaultdict(list))

        # kto hyperparams
        self.beta = finetuning_args.pref_beta
        self.desirable_weight = finetuning_args.kto_chosen_weight
        self.undesirable_weight = finetuning_args.kto_rejected_weight
        self.ftx_gamma = finetuning_args.pref_ftx
        # trl
        # Not all losses require a KL calculation
        self.calculate_KL = True
        if hasattr(self, "loss_type") and self.loss_type in ["apo_zero_unpaired"]:
            self.calculate_KL = False
        else:
            self.loss_type = "kto"

        Trainer.__init__(self, model=model, **kwargs)
        self.model_accepts_loss_kwargs = False  # overwrite trainer's default behavior
        if not hasattr(self, "accelerator"):
            raise AttributeError("Please update `transformers`.")

        warnings.simplefilter("ignore")  # remove gc warnings on ref model

        if ref_model is not None:
            if self.is_deepspeed_enabled:
                if not (
                    getattr(ref_model, "is_loaded_in_8bit", False) or getattr(ref_model, "is_loaded_in_4bit", False)
                ):  # quantized models are already set on the correct device
                    self.ref_model = self._prepare_deepspeed(self.ref_model)
            else:
                self.ref_model = self.accelerator.prepare_model(self.ref_model, evaluation_mode=True)
                self.ref_model.eval()

        if processor is not None:
            self.add_callback(SaveProcessorCallback(processor))

        if finetuning_args.use_badam:
            from badam import BAdamCallback, clip_grad_norm_old_version  # type: ignore

Review comment (refactor):
For **KTOTrainer**, it mainly relies on some functions from **TRL**, while the rest are overridden.

1. **`kto_loss`**
   This function is used directly via `self.kto_loss(...)`.

2. **`prepare_deepspeed`**
   This function is also used directly.

3. **`compute_loss`**
   It is invoked through

   ```python
   return super().compute_loss(model, inputs, return_outputs)
   ```

   Fortunately, the implementation remains unchanged, so this part is safe.

4. **`prediction_step` / `evaluation_loop`** (used by `trainer.evaluate()`)
   As long as the code runs correctly, this is not a concern and should be fine.

5. **`disable_dropout_in_model`**: luckily, it has no change.
---

### About `kto_loss`

There is no major difference between **TRL 0.9.6** and **0.24.0**, except for the introduction of the `apo_zero_unpaired` option.
This option will set `self.calculate_kl` to `False`.

To avoid introducing large behavioral changes, I suggest explicitly setting `self.calculate_kl = True`.
Otherwise, the `apo_zero_unpaired` parameter may have **system-wide effects**.

---

### About `prepare_deepspeed`

Currently, it is used as:

```python
self.ref_model = self._prepare_deepspeed(self.ref_model)
```

However, starting from **TRL ≥ 0.18.0**, `_prepare_deepspeed` is no longer an internal class method.
Instead, it should be used as a standalone utility function:

```python
from ..models import create_reference_model, prepare_deepspeed

self.ref_model = prepare_deepspeed(self.ref_model, self.accelerator)
```

---

From the **KTOTrainer** perspective, setting the minimum supported TRL version to **0.18.0** is therefore more reasonable. And as deepspeed will be LF's mainly distribued backend, it's a good idea to try both in deepspeed and fsdp to avoid missing some cases. 

I hope it could be helpful, and free feel to poinit out any mistakes! It will not be a huge change,

@UsernameFull @hiyouga 

Code after:
                disable_dropout_in_model(ref_model)

        self.finetuning_args = finetuning_args
        self.reference_free = False
        self.use_dpo_data_collator = True  # hack to avoid warning
        self.generate_during_eval = False  # disable at evaluation
        self.label_pad_token_id = IGNORE_INDEX
        self.padding_value = 0
        self.is_encoder_decoder = model.config.is_encoder_decoder
        self.precompute_ref_log_probs = False
        self._precomputed_train_ref_log_probs = False
        self._precomputed_eval_ref_log_probs = False
        self._peft_has_been_casted_to_bf16 = False

        self.ref_model = ref_model
        self._stored_metrics = defaultdict(lambda: defaultdict(list))

        # kto hyperparams
        self.beta = finetuning_args.pref_beta
        self.desirable_weight = finetuning_args.kto_chosen_weight
        self.undesirable_weight = finetuning_args.kto_rejected_weight
        self.ftx_gamma = finetuning_args.pref_ftx
        # trl
        # Not all losses require a KL calculation
        self.calculate_KL = True
        if hasattr(self, "loss_type") and self.loss_type in ["apo_zero_unpaired"]:
            self.calculate_KL = False
        else:
            self.loss_type = "kto"

        Trainer.__init__(self, model=model, **kwargs)
        self.model_accepts_loss_kwargs = False  # overwrite trainer's default behavior
        if not hasattr(self, "accelerator"):
            raise AttributeError("Please update `transformers`.")

        warnings.simplefilter("ignore")  # remove gc warnings on ref model

        if ref_model is not None:
            if self.is_deepspeed_enabled:
                if not (
                    getattr(ref_model, "is_loaded_in_8bit", False) or getattr(ref_model, "is_loaded_in_4bit", False)
                ):  # quantized models are already set on the correct device
                    self.ref_model = prepare_deepspeed(self.ref_model, self.accelerator)
            else:
                self.ref_model = self.accelerator.prepare_model(self.ref_model, evaluation_mode=True)
                self.ref_model.eval()

        if processor is not None:
            self.add_callback(SaveProcessorCallback(processor))

        if finetuning_args.use_badam: