trainer.py 13 KB
Newer Older
chenych's avatar
chenych committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
# Copyright 2024 HuggingFace Inc. and the LlamaFactory team.
#
# This code is inspired by the HuggingFace's TRL library.
# https://github.com/huggingface/trl/blob/v0.8.0/trl/trainer/kto_trainer.py
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import warnings
from collections import defaultdict
from contextlib import nullcontext
from types import MethodType
luopl's avatar
luopl committed
22
from typing import TYPE_CHECKING, Dict, List, Literal, Optional, Tuple, Union
chenych's avatar
chenych committed
23
24
25
26
27

import torch
from transformers import Trainer
from trl import KTOTrainer
from trl.trainer import disable_dropout_in_model
luopl's avatar
luopl committed
28
from typing_extensions import override
chenych's avatar
chenych committed
29
30

from ...extras.constants import IGNORE_INDEX
chenych's avatar
chenych committed
31
from ...extras.packages import is_transformers_version_greater_than
chenych's avatar
chenych committed
32
from ..callbacks import SaveProcessorCallback
luopl's avatar
luopl committed
33
from ..trainer_utils import create_custom_optimizer, create_custom_scheduler, get_batch_logps, nested_detach
chenych's avatar
chenych committed
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52


if TYPE_CHECKING:
    import torch.utils.data
    from transformers import PreTrainedModel, ProcessorMixin

    from ...hparams import FinetuningArguments


class CustomKTOTrainer(KTOTrainer):
    def __init__(
        self,
        model: Union["PreTrainedModel", torch.nn.Module],
        ref_model: Optional[Union["PreTrainedModel", torch.nn.Module]],
        finetuning_args: "FinetuningArguments",
        processor: Optional["ProcessorMixin"],
        disable_dropout: bool = True,
        **kwargs,
    ):
luopl's avatar
luopl committed
53
54
55
        if is_transformers_version_greater_than("4.46"):
            kwargs["processing_class"] = kwargs.pop("tokenizer")

chenych's avatar
chenych committed
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
        if disable_dropout:
            disable_dropout_in_model(model)
            if ref_model is not None:
                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

        Trainer.__init__(self, model=model, **kwargs)
luopl's avatar
luopl committed
83
        self.model_accepts_loss_kwargs = False  # overwrite trainer's default behavior
chenych's avatar
chenych committed
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
        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:
luopl's avatar
luopl committed
103
            from badam import BAdamCallback, clip_grad_norm_old_version  # type: ignore
chenych's avatar
chenych committed
104
105
106
107

            self.accelerator.clip_grad_norm_ = MethodType(clip_grad_norm_old_version, self.accelerator)
            self.add_callback(BAdamCallback)

luopl's avatar
luopl committed
108
    @override
chenych's avatar
chenych committed
109
110
111
112
113
    def create_optimizer(self) -> "torch.optim.Optimizer":
        if self.optimizer is None:
            self.optimizer = create_custom_optimizer(self.model, self.args, self.finetuning_args)
        return super().create_optimizer()

luopl's avatar
luopl committed
114
    @override
chenych's avatar
chenych committed
115
116
117
118
119
120
    def create_scheduler(
        self, num_training_steps: int, optimizer: Optional["torch.optim.Optimizer"] = None
    ) -> "torch.optim.lr_scheduler.LRScheduler":
        create_custom_scheduler(self.args, num_training_steps, optimizer)
        return super().create_scheduler(num_training_steps, optimizer)

luopl's avatar
luopl committed
121
    @override
chenych's avatar
chenych committed
122
123
124
125
    def _get_train_sampler(self) -> Optional["torch.utils.data.Sampler"]:
        r"""
        Replaces the sequential sampler of KTO Trainer created by trl with the random sampler.
        """
luopl's avatar
luopl committed
126
127
128
        if self.finetuning_args.disable_shuffling:
            return torch.utils.data.SequentialSampler(self.train_dataset)

chenych's avatar
chenych committed
129
130
        return Trainer._get_train_sampler(self)

luopl's avatar
luopl committed
131
132
133
134
135
136
137
    @override
    def get_batch_samples(self, epoch_iterator, num_batches):
        r"""
        Replaces the method of KTO Trainer with the one of the standard Trainer.
        """
        return Trainer.get_batch_samples(self, epoch_iterator, num_batches)

luopl's avatar
luopl committed
138
    @override
chenych's avatar
chenych committed
139
140
    def forward(
        self, model: "PreTrainedModel", batch: Dict[str, "torch.Tensor"], prefix: Literal["", "kl_"] = ""
luopl's avatar
luopl committed
141
    ) -> Tuple["torch.Tensor", "torch.Tensor", "torch.Tensor"]:
chenych's avatar
chenych committed
142
143
144
        r"""
        Runs forward pass and computes the log probabilities.
        """
luopl's avatar
luopl committed
145
        batch = nested_detach(batch, clone=True)  # avoid error
chenych's avatar
chenych committed
146
        model_inputs = {
luopl's avatar
luopl committed
147
148
            "input_ids": batch[f"{prefix}input_ids"],
            "attention_mask": batch[f"{prefix}attention_mask"],
chenych's avatar
chenych committed
149
        }
luopl's avatar
luopl committed
150
151
        if f"{prefix}token_type_ids" in batch:
            model_inputs["token_type_ids"] = batch[f"{prefix}token_type_ids"]
luopl's avatar
luopl committed
152

chenych's avatar
chenych committed
153
154
155
        if "pixel_values" in batch:
            model_inputs["pixel_values"] = batch["pixel_values"]

luopl's avatar
luopl committed
156
157
        if "image_grid_thw" in batch:
            model_inputs["image_grid_thw"] = batch["image_grid_thw"]
chenych's avatar
chenych committed
158

chenych's avatar
chenych committed
159
160
161
162
163
164
165
166
167
        if "aspect_ratio_ids" in batch:
            model_inputs["aspect_ratio_ids"] = batch["aspect_ratio_ids"]

        if "aspect_ratio_mask" in batch:
            model_inputs["aspect_ratio_mask"] = batch["aspect_ratio_mask"]

        if f"{prefix}cross_attention_mask" in batch:
            model_inputs["cross_attention_mask"] = batch[f"{prefix}cross_attention_mask"]

chenych's avatar
chenych committed
168
        logits = model(**model_inputs, return_dict=True, use_cache=False).logits.to(torch.float32)
luopl's avatar
luopl committed
169
170
        logps, valid_length = get_batch_logps(logits=logits, labels=batch[f"{prefix}labels"])
        return logits, logps, logps / valid_length
chenych's avatar
chenych committed
171

luopl's avatar
luopl committed
172
    @override
chenych's avatar
chenych committed
173
174
    def concatenated_forward(
        self, model: "PreTrainedModel", batch: Dict[str, "torch.Tensor"]
luopl's avatar
luopl committed
175
176
    ) -> Tuple["torch.Tensor", "torch.Tensor", "torch.Tensor", "torch.Tensor", "torch.Tensor", "torch.Tensor"]:
        target_logits, target_logps, target_logps_avg = self.forward(model, batch)
chenych's avatar
chenych committed
177
        with torch.no_grad():
luopl's avatar
luopl committed
178
            _, kl_logps, _ = self.forward(model, batch, prefix="kl_")
chenych's avatar
chenych committed
179
180
181
182

        if len(target_logps) != len(batch["kto_tags"]):
            raise ValueError("Mismatched shape of inputs and labels.")

luopl's avatar
luopl committed
183
        chosen_logits = target_logits[batch["kto_tags"]]
chenych's avatar
chenych committed
184
        chosen_logps = target_logps[batch["kto_tags"]]
luopl's avatar
luopl committed
185
        rejected_logits = target_logits[~batch["kto_tags"]]
chenych's avatar
chenych committed
186
187
        rejected_logps = target_logps[~batch["kto_tags"]]
        chosen_logps_avg = target_logps_avg[batch["kto_tags"]]
luopl's avatar
luopl committed
188
        return chosen_logps, rejected_logps, chosen_logits, rejected_logits, kl_logps, chosen_logps_avg
chenych's avatar
chenych committed
189

luopl's avatar
luopl committed
190
    @override
chenych's avatar
chenych committed
191
192
193
194
195
196
197
198
199
200
201
202
203
204
    def compute_reference_log_probs(
        self, model: "PreTrainedModel", batch: Dict[str, "torch.Tensor"]
    ) -> Tuple["torch.Tensor", "torch.Tensor", "torch.Tensor"]:
        r"""
        Computes log probabilities of the reference model.
        """
        if self.ref_model is None:
            ref_model = model
            ref_context = self.accelerator.unwrap_model(model).disable_adapter()
        else:
            ref_model = self.ref_model
            ref_context = nullcontext()

        with torch.no_grad(), ref_context:
luopl's avatar
luopl committed
205
            reference_chosen_logps, reference_rejected_logps, _, _, reference_kl_logps, _ = self.concatenated_forward(
chenych's avatar
chenych committed
206
207
208
209
210
                ref_model, batch
            )

        return reference_chosen_logps, reference_rejected_logps, reference_kl_logps

luopl's avatar
luopl committed
211
    @override
chenych's avatar
chenych committed
212
213
214
215
216
217
218
219
220
    def get_batch_loss_metrics(
        self,
        model: "PreTrainedModel",
        batch: Dict[str, "torch.Tensor"],
    ) -> Tuple["torch.Tensor", Dict[str, "torch.Tensor"]]:
        r"""
        Computes the DPO loss and other metrics for the given batch of inputs for train or test.
        """
        metrics = {}
luopl's avatar
luopl committed
221
222
223
224
225
226
227
228
        (
            policy_chosen_logps,
            policy_rejected_logps,
            policy_chosen_logits,
            policy_rejected_logits,
            policy_kl_logps,
            policy_chosen_logps_avg,
        ) = self.concatenated_forward(model, batch)
chenych's avatar
chenych committed
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
        reference_chosen_logps, reference_rejected_logps, reference_kl_logps = self.compute_reference_log_probs(
            model, batch
        )
        losses, chosen_rewards, rejected_rewards, kl = self.kto_loss(
            policy_chosen_logps,
            policy_rejected_logps,
            policy_kl_logps,
            reference_chosen_logps,
            reference_rejected_logps,
            reference_kl_logps,
        )
        losses = losses.nanmean()

        if self.ftx_gamma > 1e-6 and len(policy_chosen_logps) > 0:  # remember to rescale
            sft_loss = -policy_chosen_logps_avg
            losses += self.ftx_gamma * sft_loss.nanmean() / len(policy_chosen_logps) * len(batch["labels"])

luopl's avatar
luopl committed
246
247
248
249
250
251
252
        num_chosen = len(chosen_rewards)
        num_rejected = len(rejected_rewards)
        if num_chosen > 0:
            metrics["rewards/chosen_sum"] = chosen_rewards.nansum().item()
            metrics["logps/chosen_sum"] = policy_chosen_logps.nansum().item()
            metrics["logits/chosen_sum"] = policy_chosen_logits.nansum().item()
            metrics["count/chosen"] = float(num_chosen)
chenych's avatar
chenych committed
253

luopl's avatar
luopl committed
254
255
256
257
258
        if num_rejected > 0:
            metrics["rewards/rejected_sum"] = rejected_rewards.nansum().item()
            metrics["logps/rejected_sum"] = policy_rejected_logps.nansum().item()
            metrics["logits/rejected_sum"] = policy_rejected_logits.nansum().item()
            metrics["count/rejected"] = float(num_rejected)
chenych's avatar
chenych committed
259

luopl's avatar
luopl committed
260
261
        metrics["kl"] = kl.item()
        return losses, metrics
chenych's avatar
chenych committed
262

luopl's avatar
luopl committed
263
    @override
luopl's avatar
luopl committed
264
265
266
    def compute_loss(
        self, model: "PreTrainedModel", inputs: Dict[str, "torch.Tensor"], return_outputs: bool = False, **kwargs
    ) -> Union["torch.Tensor", Tuple["torch.Tensor", List["torch.Tensor"]]]:
luopl's avatar
luopl committed
267
        r"""
chenych's avatar
chenych committed
268
        Subclass and override to accept extra kwargs.
luopl's avatar
luopl committed
269
        """
chenych's avatar
chenych committed
270
        return super().compute_loss(model, inputs, return_outputs)
chenych's avatar
chenych committed
271

luopl's avatar
luopl committed
272
    @override
chenych's avatar
chenych committed
273
    def log(self, logs: Dict[str, float], *args, **kwargs) -> None:
luopl's avatar
luopl committed
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
        r"""
        Log `logs` on the various objects watching training, including stored metrics.
        """
        # logs either has "loss" or "eval_loss"
        train_eval = "train" if "loss" in logs else "eval"
        prefix = "eval_" if train_eval == "eval" else ""
        # Add averaged stored metrics to logs
        key_list, metric_list = [], []
        for key, metrics in self._stored_metrics[train_eval].items():
            key_list.append(key)
            metric_list.append(torch.tensor(metrics, dtype=torch.float).to(self.accelerator.device).sum().item())

        del self._stored_metrics[train_eval]
        if len(metric_list) < 9:  # pad to for all reduce
            for i in range(9 - len(metric_list)):
                key_list.append(f"dummy_{i}")
                metric_list.append(0.0)

        metric_list = torch.tensor(metric_list, dtype=torch.float).to(self.accelerator.device)
        metric_list = self.accelerator.reduce(metric_list, "sum").tolist()
        metric_dict: Dict[str, float] = dict(zip(key_list, metric_list))
        for split in ["chosen", "rejected"]:  # accumulate average metrics from sums and lengths
            if f"count/{split}" in metric_dict:
                for key in ("rewards", "logps", "logits"):
                    logs[f"{prefix}{key}/{split}"] = metric_dict[f"{key}/{split}_sum"] / metric_dict[f"count/{split}"]
                    del metric_dict[f"{key}/{split}_sum"]
                del metric_dict[f"count/{split}"]

        if f"{prefix}rewards/chosen" in logs and f"{prefix}rewards/rejected" in logs:  # calculate reward margin
            logs[f"{prefix}rewards/margins"] = logs[f"{prefix}rewards/chosen"] - logs[f"{prefix}rewards/rejected"]

        for key, metric in metric_dict.items():  # add remaining items
            if not key.startswith("dummy_"):
                logs[key] = metric

chenych's avatar
chenych committed
309
        return Trainer.log(self, logs, *args, **kwargs)