trainer_utils.py 27.6 KB
Newer Older
chenych's avatar
chenych committed
1
# Copyright 2025 HuggingFace Inc. and the LlamaFactory team.
chenych's avatar
chenych committed
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
#
# This code is inspired by the original GaLore's implementation: https://github.com/jiaweizzhao/GaLore
# and the original LoRA+'s implementation: https://github.com/nikhil-ghosh-berkeley/loraplus
# and the original BAdam's implementation: https://github.com/Ledzy/BAdam
# and the HuggingFace's TRL library: https://github.com/huggingface/trl
#
# 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.

chenych's avatar
chenych committed
20
21
import json
import os
luopl's avatar
luopl committed
22
23
from collections.abc import Mapping
from pathlib import Path
chenych's avatar
chenych committed
24
from typing import TYPE_CHECKING, Any, Callable, Optional, Union
Rayyyyy's avatar
V0.6.3  
Rayyyyy committed
25
26
27

import torch
from transformers import Trainer
chenych's avatar
chenych committed
28
29
from transformers.integrations import is_deepspeed_zero3_enabled
from transformers.modeling_utils import is_fsdp_enabled
Rayyyyy's avatar
V0.6.3  
Rayyyyy committed
30
31
32
from transformers.optimization import get_scheduler
from transformers.pytorch_utils import ALL_LAYERNORM_LAYERS
from transformers.trainer_pt_utils import get_parameter_names
luopl's avatar
luopl committed
33
from typing_extensions import override
Rayyyyy's avatar
V0.6.3  
Rayyyyy committed
34

luopl's avatar
luopl committed
35
from ..extras import logging
chenych's avatar
chenych committed
36
from ..extras.constants import IGNORE_INDEX, SWANLAB_CONFIG
luopl's avatar
luopl committed
37
from ..extras.packages import is_apollo_available, is_galore_available, is_ray_available
Rayyyyy's avatar
V0.6.3  
Rayyyyy committed
38
39
40
41
42
from ..hparams import FinetuningArguments, ModelArguments
from ..model import find_all_linear_modules, load_model, load_tokenizer, load_valuehead_params


if is_galore_available():
luopl's avatar
luopl committed
43
44
45
46
47
48
49
50
51
52
    from galore_torch import GaLoreAdafactor, GaLoreAdamW, GaLoreAdamW8bit  # type: ignore


if is_apollo_available():
    from apollo_torch import APOLLOAdamW  # type: ignore


if is_ray_available():
    from ray.train import RunConfig, ScalingConfig
    from ray.train.torch import TorchTrainer
Rayyyyy's avatar
V0.6.3  
Rayyyyy committed
53
54
55


if TYPE_CHECKING:
chenych's avatar
chenych committed
56
    from transformers import PreTrainedModel, TrainerCallback, TrainerState
Rayyyyy's avatar
V0.6.3  
Rayyyyy committed
57
58
    from trl import AutoModelForCausalLMWithValueHead

luopl's avatar
luopl committed
59
    from ..hparams import DataArguments, RayArguments, TrainingArguments
Rayyyyy's avatar
V0.6.3  
Rayyyyy committed
60
61


luopl's avatar
luopl committed
62
logger = logging.get_logger(__name__)
Rayyyyy's avatar
V0.6.3  
Rayyyyy committed
63
64
65


class DummyOptimizer(torch.optim.Optimizer):
chenych's avatar
chenych committed
66
    r"""A dummy optimizer used for the GaLore or APOLLO algorithm."""
Rayyyyy's avatar
V0.6.3  
Rayyyyy committed
67
68

    def __init__(
chenych's avatar
chenych committed
69
        self, lr: float = 1e-3, optimizer_dict: Optional[dict["torch.nn.Parameter", "torch.optim.Optimizer"]] = None
Rayyyyy's avatar
V0.6.3  
Rayyyyy committed
70
71
72
73
74
    ) -> None:
        dummy_tensor = torch.randn(1, 1)
        self.optimizer_dict = optimizer_dict
        super().__init__([dummy_tensor], {"lr": lr})

luopl's avatar
luopl committed
75
    @override
Rayyyyy's avatar
V0.6.3  
Rayyyyy committed
76
77
78
    def zero_grad(self, set_to_none: bool = True) -> None:
        pass

luopl's avatar
luopl committed
79
    @override
Rayyyyy's avatar
V0.6.3  
Rayyyyy committed
80
81
82
83
84
85
86
87
    def step(self, closure: Optional[Callable[[], float]] = None) -> Optional[float]:
        pass


def create_modelcard_and_push(
    trainer: "Trainer",
    model_args: "ModelArguments",
    data_args: "DataArguments",
luopl's avatar
luopl committed
88
    training_args: "TrainingArguments",
Rayyyyy's avatar
V0.6.3  
Rayyyyy committed
89
90
91
92
93
94
95
96
    finetuning_args: "FinetuningArguments",
) -> None:
    kwargs = {
        "tasks": "text-generation",
        "finetuned_from": model_args.model_name_or_path,
        "tags": ["llama-factory", finetuning_args.finetuning_type],
    }
    if data_args.dataset is not None:
chenych's avatar
chenych committed
97
98
99
100
        kwargs["dataset"] = data_args.dataset

    if model_args.use_unsloth:
        kwargs["tags"] = kwargs["tags"] + ["unsloth"]
Rayyyyy's avatar
V0.6.3  
Rayyyyy committed
101
102
103
104
105
106
107
108
109
110
111
112

    if not training_args.do_train:
        pass
    elif training_args.push_to_hub:
        trainer.push_to_hub(**kwargs)
    else:
        trainer.create_model_card(license="other", **kwargs)  # prevent from connecting to hub


def create_ref_model(
    model_args: "ModelArguments", finetuning_args: "FinetuningArguments", add_valuehead: bool = False
) -> Optional[Union["PreTrainedModel", "AutoModelForCausalLMWithValueHead"]]:
chenych's avatar
chenych committed
113
    r"""Create reference model for PPO/DPO training. Evaluation mode is not supported.
Rayyyyy's avatar
V0.6.3  
Rayyyyy committed
114
115
116
117

    The valuehead parameter is randomly initialized since it is useless for PPO training.
    """
    if finetuning_args.ref_model is not None:
chenych's avatar
chenych committed
118
119
120
121
122
        ref_model_args = ModelArguments.copyfrom(
            model_args,
            model_name_or_path=finetuning_args.ref_model,
            adapter_name_or_path=finetuning_args.ref_model_adapters,
            quantization_bit=finetuning_args.ref_model_quantization_bit,
Rayyyyy's avatar
V0.6.3  
Rayyyyy committed
123
        )
chenych's avatar
chenych committed
124
125
        ref_finetuning_args = FinetuningArguments()
        tokenizer = load_tokenizer(ref_model_args)["tokenizer"]
Rayyyyy's avatar
V0.6.3  
Rayyyyy committed
126
127
128
        ref_model = load_model(
            tokenizer, ref_model_args, ref_finetuning_args, is_trainable=False, add_valuehead=add_valuehead
        )
luopl's avatar
luopl committed
129
        logger.info_rank0(f"Created reference model from {finetuning_args.ref_model}")
Rayyyyy's avatar
V0.6.3  
Rayyyyy committed
130
131
132
133
    else:
        if finetuning_args.finetuning_type == "lora":
            ref_model = None
        else:
chenych's avatar
chenych committed
134
135
136
            ref_model_args = ModelArguments.copyfrom(model_args)
            ref_finetuning_args = FinetuningArguments()
            tokenizer = load_tokenizer(ref_model_args)["tokenizer"]
Rayyyyy's avatar
V0.6.3  
Rayyyyy committed
137
            ref_model = load_model(
chenych's avatar
chenych committed
138
                tokenizer, ref_model_args, ref_finetuning_args, is_trainable=False, add_valuehead=add_valuehead
Rayyyyy's avatar
V0.6.3  
Rayyyyy committed
139
            )
luopl's avatar
luopl committed
140
            logger.info_rank0("Created reference model from the model itself.")
Rayyyyy's avatar
V0.6.3  
Rayyyyy committed
141
142
143
144
145
146
147

    return ref_model


def create_reward_model(
    model: "AutoModelForCausalLMWithValueHead", model_args: "ModelArguments", finetuning_args: "FinetuningArguments"
) -> Optional["AutoModelForCausalLMWithValueHead"]:
chenych's avatar
chenych committed
148
    r"""Create reward model for PPO training."""
Rayyyyy's avatar
V0.6.3  
Rayyyyy committed
149
150
    if finetuning_args.reward_model_type == "api":
        assert finetuning_args.reward_model.startswith("http"), "Please provide full url."
luopl's avatar
luopl committed
151
        logger.info_rank0(f"Use reward server {finetuning_args.reward_model}")
Rayyyyy's avatar
V0.6.3  
Rayyyyy committed
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
        return finetuning_args.reward_model
    elif finetuning_args.reward_model_type == "lora":
        model.pretrained_model.load_adapter(finetuning_args.reward_model, "reward")
        for name, param in model.named_parameters():  # https://github.com/huggingface/peft/issues/1090
            if "default" in name:
                param.data = param.data.to(torch.float32)  # trainable params should in fp32
        vhead_params = load_valuehead_params(finetuning_args.reward_model, model_args)
        assert vhead_params is not None, "Reward model is not correctly loaded."
        model.register_buffer("reward_head_weight", vhead_params["v_head.summary.weight"], persistent=False)
        model.register_buffer("reward_head_bias", vhead_params["v_head.summary.bias"], persistent=False)
        model.register_buffer(
            "default_head_weight", torch.zeros_like(vhead_params["v_head.summary.weight"]), persistent=False
        )
        model.register_buffer(
            "default_head_bias", torch.zeros_like(vhead_params["v_head.summary.bias"]), persistent=False
        )
luopl's avatar
luopl committed
168
        logger.info_rank0(f"Loaded adapter weights of reward model from {finetuning_args.reward_model}")
Rayyyyy's avatar
V0.6.3  
Rayyyyy committed
169
170
        return None
    else:
chenych's avatar
chenych committed
171
172
173
174
175
        reward_model_args = ModelArguments.copyfrom(
            model_args,
            model_name_or_path=finetuning_args.reward_model,
            adapter_name_or_path=finetuning_args.reward_model_adapters,
            quantization_bit=finetuning_args.reward_model_quantization_bit,
Rayyyyy's avatar
V0.6.3  
Rayyyyy committed
176
        )
chenych's avatar
chenych committed
177
178
        reward_finetuning_args = FinetuningArguments()
        tokenizer = load_tokenizer(reward_model_args)["tokenizer"]
Rayyyyy's avatar
V0.6.3  
Rayyyyy committed
179
180
181
        reward_model = load_model(
            tokenizer, reward_model_args, reward_finetuning_args, is_trainable=False, add_valuehead=True
        )
luopl's avatar
luopl committed
182
183
        logger.info_rank0(f"Loaded full weights of reward model from {finetuning_args.reward_model}")
        logger.warning_rank0("Please ensure the ppo model and reward model share SAME tokenizer and vocabulary.")
Rayyyyy's avatar
V0.6.3  
Rayyyyy committed
184
185
186
        return reward_model


chenych's avatar
chenych committed
187
188
def _get_decay_parameter_names(model: "PreTrainedModel") -> list[str]:
    r"""Return a list of names of parameters with weight decay. (weights in non-layernorm layers)."""
Rayyyyy's avatar
V0.6.3  
Rayyyyy committed
189
190
191
192
193
194
195
    decay_parameters = get_parameter_names(model, ALL_LAYERNORM_LAYERS)
    decay_parameters = [name for name in decay_parameters if "bias" not in name]
    return decay_parameters


def _create_galore_optimizer(
    model: "PreTrainedModel",
luopl's avatar
luopl committed
196
    training_args: "TrainingArguments",
Rayyyyy's avatar
V0.6.3  
Rayyyyy committed
197
198
199
    finetuning_args: "FinetuningArguments",
) -> "torch.optim.Optimizer":
    if len(finetuning_args.galore_target) == 1 and finetuning_args.galore_target[0] == "all":
chenych's avatar
chenych committed
200
        galore_targets = find_all_linear_modules(model, finetuning_args.freeze_vision_tower)
Rayyyyy's avatar
V0.6.3  
Rayyyyy committed
201
202
203
    else:
        galore_targets = finetuning_args.galore_target

chenych's avatar
chenych committed
204
    galore_params: list[torch.nn.Parameter] = []
Rayyyyy's avatar
V0.6.3  
Rayyyyy committed
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
    for name, module in model.named_modules():
        if isinstance(module, torch.nn.Linear) and any(target in name for target in galore_targets):
            for param in module.parameters():
                if param.requires_grad and len(param.shape) > 1:
                    galore_params.append(param)

    galore_kwargs = {
        "rank": finetuning_args.galore_rank,
        "update_proj_gap": finetuning_args.galore_update_interval,
        "scale": finetuning_args.galore_scale,
        "proj_type": finetuning_args.galore_proj_type,
    }

    id_galore_params = {id(param) for param in galore_params}
    decay_params, nodecay_params = [], []  # they are non-galore parameters
chenych's avatar
chenych committed
220
    trainable_params: list[torch.nn.Parameter] = []  # galore_params + decay_params + nodecay_params
Rayyyyy's avatar
V0.6.3  
Rayyyyy committed
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
    decay_param_names = _get_decay_parameter_names(model)
    for name, param in model.named_parameters():
        if param.requires_grad:
            trainable_params.append(param)
            if id(param) not in id_galore_params:
                if name in decay_param_names:
                    decay_params.append(param)
                else:
                    nodecay_params.append(param)

    _, optim_kwargs = Trainer.get_optimizer_cls_and_kwargs(training_args)

    if training_args.optim == "adamw_torch":
        optim_class = GaLoreAdamW
    elif training_args.optim in ["adamw_bnb_8bit", "adamw_8bit", "paged_adamw_8bit"]:
        optim_class = GaLoreAdamW8bit
    elif training_args.optim == "adafactor":
        optim_class = GaLoreAdafactor
    else:
luopl's avatar
luopl committed
240
        raise NotImplementedError(f"Unknown optim: {training_args.optim}.")
Rayyyyy's avatar
V0.6.3  
Rayyyyy committed
241
242

    if finetuning_args.galore_layerwise:
luopl's avatar
luopl committed
243
        logger.warning_rank0("The displayed gradient norm will be all zeros in layerwise GaLore.")
Rayyyyy's avatar
V0.6.3  
Rayyyyy committed
244
245
246
        if training_args.gradient_accumulation_steps != 1:
            raise ValueError("Per-layer GaLore does not support gradient accumulation.")

chenych's avatar
chenych committed
247
        optimizer_dict: dict[torch.Tensor, torch.optim.Optimizer] = {}
Rayyyyy's avatar
V0.6.3  
Rayyyyy committed
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
        for param in nodecay_params:
            param_groups = [dict(params=[param], weight_decay=0.0)]
            optimizer_dict[param] = optim_class(param_groups, **optim_kwargs)
        for param in decay_params:
            param_groups = [dict(params=[param], weight_decay=training_args.weight_decay)]
            optimizer_dict[param] = optim_class(param_groups, **optim_kwargs)
        for param in galore_params:  # galore params have weight decay
            param_groups = [dict(params=[param], weight_decay=training_args.weight_decay, **galore_kwargs)]
            optimizer_dict[param] = optim_class(param_groups, **optim_kwargs)

        def optimizer_hook(param: "torch.nn.Parameter"):
            if param.grad is not None:
                optimizer_dict[param].step()
                optimizer_dict[param].zero_grad()

        for param in trainable_params:
            param.register_post_accumulate_grad_hook(optimizer_hook)

        optimizer = DummyOptimizer(lr=training_args.learning_rate, optimizer_dict=optimizer_dict)
    else:
        param_groups = [
            dict(params=nodecay_params, weight_decay=0.0),
            dict(params=decay_params, weight_decay=training_args.weight_decay),
            dict(params=galore_params, weight_decay=training_args.weight_decay, **galore_kwargs),
        ]
        optimizer = optim_class(param_groups, **optim_kwargs)

luopl's avatar
luopl committed
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
    logger.info_rank0(
        f"Using GaLore optimizer with args: {galore_kwargs}. "
        "It may cause hanging at the start of training, wait patiently."
    )
    return optimizer


def _create_apollo_optimizer(
    model: "PreTrainedModel",
    training_args: "TrainingArguments",
    finetuning_args: "FinetuningArguments",
) -> "torch.optim.Optimizer":
    if len(finetuning_args.apollo_target) == 1 and finetuning_args.apollo_target[0] == "all":
        apollo_targets = find_all_linear_modules(model, finetuning_args.freeze_vision_tower)
    else:
        apollo_targets = finetuning_args.apollo_target

chenych's avatar
chenych committed
292
    apollo_params: list[torch.nn.Parameter] = []
luopl's avatar
luopl committed
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
    for name, module in model.named_modules():
        if isinstance(module, torch.nn.Linear) and any(target in name for target in apollo_targets):
            for param in module.parameters():
                if param.requires_grad and len(param.shape) > 1:
                    apollo_params.append(param)

    apollo_kwargs = {
        "rank": finetuning_args.apollo_rank,
        "proj": finetuning_args.apollo_proj,
        "proj_type": finetuning_args.apollo_proj_type,
        "update_proj_gap": finetuning_args.apollo_update_interval,
        "scale": finetuning_args.apollo_scale,
        "scale_type": finetuning_args.apollo_scale_type,
        "scale_front": finetuning_args.apollo_scale_front,
    }

    id_apollo_params = {id(param) for param in apollo_params}
    decay_params, nodecay_params = [], []  # they are non-apollo parameters
chenych's avatar
chenych committed
311
    trainable_params: list[torch.nn.Parameter] = []  # apollo_params + decay_params + nodecay_params
luopl's avatar
luopl committed
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
    decay_param_names = _get_decay_parameter_names(model)
    for name, param in model.named_parameters():
        if param.requires_grad:
            trainable_params.append(param)
            if id(param) not in id_apollo_params:
                if name in decay_param_names:
                    decay_params.append(param)
                else:
                    nodecay_params.append(param)

    _, optim_kwargs = Trainer.get_optimizer_cls_and_kwargs(training_args)

    if training_args.optim == "adamw_torch":
        optim_class = APOLLOAdamW
    else:
        raise NotImplementedError(f"Unknown optim: {training_args.optim}.")

    if finetuning_args.apollo_layerwise:
        logger.warning_rank0("The displayed gradient norm will be all zeros in layerwise APOLLO.")
        if training_args.gradient_accumulation_steps != 1:
            raise ValueError("Per-layer APOLLO does not support gradient accumulation.")

chenych's avatar
chenych committed
334
        optimizer_dict: dict[torch.Tensor, torch.optim.Optimizer] = {}
luopl's avatar
luopl committed
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
        for param in nodecay_params:
            param_groups = [dict(params=[param], weight_decay=0.0)]
            optimizer_dict[param] = optim_class(param_groups, **optim_kwargs)
        for param in decay_params:
            param_groups = [dict(params=[param], weight_decay=training_args.weight_decay)]
            optimizer_dict[param] = optim_class(param_groups, **optim_kwargs)
        for param in apollo_params:  # apollo params have weight decay
            param_groups = [dict(params=[param], weight_decay=training_args.weight_decay, **apollo_kwargs)]
            optimizer_dict[param] = optim_class(param_groups, **optim_kwargs)

        def optimizer_hook(param: "torch.nn.Parameter"):
            if param.grad is not None:
                optimizer_dict[param].step()
                optimizer_dict[param].zero_grad()

        for param in trainable_params:
            param.register_post_accumulate_grad_hook(optimizer_hook)

        optimizer = DummyOptimizer(lr=training_args.learning_rate, optimizer_dict=optimizer_dict)
    else:
        param_groups = [
            dict(params=nodecay_params, weight_decay=0.0),
            dict(params=decay_params, weight_decay=training_args.weight_decay),
            dict(params=apollo_params, weight_decay=training_args.weight_decay, **apollo_kwargs),
        ]
        optimizer = optim_class(param_groups, **optim_kwargs)

    logger.info_rank0(f"Using APOLLO optimizer with args: {apollo_kwargs}.")
Rayyyyy's avatar
V0.6.3  
Rayyyyy committed
363
364
365
366
367
    return optimizer


def _create_loraplus_optimizer(
    model: "PreTrainedModel",
luopl's avatar
luopl committed
368
    training_args: "TrainingArguments",
Rayyyyy's avatar
V0.6.3  
Rayyyyy committed
369
370
371
372
373
374
375
    finetuning_args: "FinetuningArguments",
) -> "torch.optim.Optimizer":
    default_lr = training_args.learning_rate
    loraplus_lr = training_args.learning_rate * finetuning_args.loraplus_lr_ratio
    embedding_lr = finetuning_args.loraplus_lr_embedding

    decay_param_names = _get_decay_parameter_names(model)
chenych's avatar
chenych committed
376
    param_dict: dict[str, list[torch.nn.Parameter]] = {
Rayyyyy's avatar
V0.6.3  
Rayyyyy committed
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
        "lora_a": [],
        "lora_b": [],
        "lora_b_nodecay": [],
        "embedding": [],
    }
    for name, param in model.named_parameters():
        if param.requires_grad:
            if "lora_embedding_B" in name:
                param_dict["embedding"].append(param)
            elif "lora_B" in name or param.ndim == 1:
                if name in decay_param_names:
                    param_dict["lora_b"].append(param)
                else:
                    param_dict["lora_b_nodecay"].append(param)
            else:
                param_dict["lora_a"].append(param)

    optim_class, optim_kwargs = Trainer.get_optimizer_cls_and_kwargs(training_args)
    param_groups = [
        dict(params=param_dict["lora_a"], lr=default_lr, weight_decay=training_args.weight_decay),
        dict(params=param_dict["lora_b"], lr=loraplus_lr, weight_decay=training_args.weight_decay),
        dict(params=param_dict["lora_b_nodecay"], lr=loraplus_lr, weight_decay=0.0),
        dict(params=param_dict["embedding"], lr=embedding_lr, weight_decay=training_args.weight_decay),
    ]
    optimizer = optim_class(param_groups, **optim_kwargs)
luopl's avatar
luopl committed
402
    logger.info_rank0(f"Using LoRA+ optimizer with loraplus lr ratio {finetuning_args.loraplus_lr_ratio:.2f}.")
Rayyyyy's avatar
V0.6.3  
Rayyyyy committed
403
404
405
406
407
    return optimizer


def _create_badam_optimizer(
    model: "PreTrainedModel",
luopl's avatar
luopl committed
408
    training_args: "TrainingArguments",
Rayyyyy's avatar
V0.6.3  
Rayyyyy committed
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
    finetuning_args: "FinetuningArguments",
) -> "torch.optim.Optimizer":
    decay_params, nodecay_params = [], []
    decay_param_names = _get_decay_parameter_names(model)
    for name, param in model.named_parameters():
        if param.requires_grad:
            if name in decay_param_names:
                decay_params.append(param)
            else:
                nodecay_params.append(param)

    optim_class, optim_kwargs = Trainer.get_optimizer_cls_and_kwargs(training_args)
    param_groups = [
        dict(params=nodecay_params, weight_decay=0.0),
        dict(params=decay_params, weight_decay=training_args.weight_decay),
    ]

    if finetuning_args.badam_mode == "layer":
luopl's avatar
luopl committed
427
        from badam import BlockOptimizer  # type: ignore
Rayyyyy's avatar
V0.6.3  
Rayyyyy committed
428
429
430
431
432
433

        base_optimizer = optim_class(param_groups, **optim_kwargs)
        optimizer = BlockOptimizer(
            base_optimizer=base_optimizer,
            named_parameters_list=list(model.named_parameters()),
            block_prefix_list=None,
chenych's avatar
chenych committed
434
            switch_block_every=finetuning_args.badam_switch_interval,
Rayyyyy's avatar
V0.6.3  
Rayyyyy committed
435
436
437
            start_block=finetuning_args.badam_start_block,
            switch_mode=finetuning_args.badam_switch_mode,
            verbose=finetuning_args.badam_verbose,
chenych's avatar
chenych committed
438
            ds_zero3_enabled=is_deepspeed_zero3_enabled(),
Rayyyyy's avatar
V0.6.3  
Rayyyyy committed
439
        )
luopl's avatar
luopl committed
440
        logger.info_rank0(
Rayyyyy's avatar
V0.6.3  
Rayyyyy committed
441
            f"Using BAdam optimizer with layer-wise update, switch mode is {finetuning_args.badam_switch_mode}, "
chenych's avatar
chenych committed
442
            f"switch block every {finetuning_args.badam_switch_interval} steps, "
Rayyyyy's avatar
V0.6.3  
Rayyyyy committed
443
444
445
446
            f"default start block is {finetuning_args.badam_start_block}"
        )

    elif finetuning_args.badam_mode == "ratio":
luopl's avatar
luopl committed
447
        from badam import BlockOptimizerRatio  # type: ignore
Rayyyyy's avatar
V0.6.3  
Rayyyyy committed
448
449
450
451
452
453
454
455
456
457
458

        assert finetuning_args.badam_update_ratio > 1e-6
        optimizer = BlockOptimizerRatio(
            param_groups=param_groups,
            named_parameters_list=list(model.named_parameters()),
            update_ratio=finetuning_args.badam_update_ratio,
            mask_mode=finetuning_args.badam_mask_mode,
            verbose=finetuning_args.badam_verbose,
            include_embedding=False,
            **optim_kwargs,
        )
luopl's avatar
luopl committed
459
        logger.info_rank0(
chenych's avatar
chenych committed
460
            f"Using BAdam optimizer with ratio-based update, update ratio is {finetuning_args.badam_update_ratio}, "
Rayyyyy's avatar
V0.6.3  
Rayyyyy committed
461
462
463
464
465
466
            f"mask mode is {finetuning_args.badam_mask_mode}"
        )

    return optimizer


chenych's avatar
chenych committed
467
468
def _create_adam_mini_optimizer(
    model: "PreTrainedModel",
luopl's avatar
luopl committed
469
    training_args: "TrainingArguments",
chenych's avatar
chenych committed
470
) -> "torch.optim.Optimizer":
luopl's avatar
luopl committed
471
    from adam_mini import Adam_mini  # type: ignore
chenych's avatar
chenych committed
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487

    hidden_size = getattr(model.config, "hidden_size", None)
    num_q_head = getattr(model.config, "num_attention_heads", None)
    num_kv_head = getattr(model.config, "num_key_value_heads", None)

    optimizer = Adam_mini(
        named_parameters=model.named_parameters(),
        lr=training_args.learning_rate,
        betas=(training_args.adam_beta1, training_args.adam_beta2),
        eps=training_args.adam_epsilon,
        weight_decay=training_args.weight_decay,
        model_sharding=is_fsdp_enabled() or is_deepspeed_zero3_enabled(),
        dim=hidden_size,
        n_heads=num_q_head,
        n_kv_heads=num_kv_head,
    )
luopl's avatar
luopl committed
488
    logger.info_rank0("Using Adam-mini optimizer.")
chenych's avatar
chenych committed
489
490
491
492
    return optimizer


def create_custom_optimizer(
Rayyyyy's avatar
V0.6.3  
Rayyyyy committed
493
    model: "PreTrainedModel",
luopl's avatar
luopl committed
494
    training_args: "TrainingArguments",
Rayyyyy's avatar
V0.6.3  
Rayyyyy committed
495
496
497
498
499
    finetuning_args: "FinetuningArguments",
) -> Optional["torch.optim.Optimizer"]:
    if finetuning_args.use_galore:
        return _create_galore_optimizer(model, training_args, finetuning_args)

luopl's avatar
luopl committed
500
501
502
    if finetuning_args.use_apollo:
        return _create_apollo_optimizer(model, training_args, finetuning_args)

Rayyyyy's avatar
V0.6.3  
Rayyyyy committed
503
504
505
506
507
508
    if finetuning_args.loraplus_lr_ratio is not None:
        return _create_loraplus_optimizer(model, training_args, finetuning_args)

    if finetuning_args.use_badam:
        return _create_badam_optimizer(model, training_args, finetuning_args)

chenych's avatar
chenych committed
509
510
511
    if finetuning_args.use_adam_mini:
        return _create_adam_mini_optimizer(model, training_args)

Rayyyyy's avatar
V0.6.3  
Rayyyyy committed
512
513

def create_custom_scheduler(
luopl's avatar
luopl committed
514
    training_args: "TrainingArguments",
Rayyyyy's avatar
V0.6.3  
Rayyyyy committed
515
516
517
    num_training_steps: int,
    optimizer: Optional["torch.optim.Optimizer"] = None,
) -> None:
chenych's avatar
chenych committed
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
    if training_args.lr_scheduler_type == "warmup_stable_decay":
        num_warmup_steps = training_args.get_warmup_steps(num_training_steps)
        remaining_steps = num_training_steps - num_warmup_steps
        num_stable_steps = remaining_steps // 3  # use 1/3 for stable by default
        num_decay_steps = remaining_steps - num_stable_steps
        scheduler_kwargs = training_args.lr_scheduler_kwargs or {}
        default_kwargs = {
            "num_stable_steps": num_stable_steps,
            "num_decay_steps": num_decay_steps,
        }
        for key, value in default_kwargs.items():
            if key not in scheduler_kwargs:
                scheduler_kwargs[key] = value

        training_args.lr_scheduler_kwargs = scheduler_kwargs

Rayyyyy's avatar
V0.6.3  
Rayyyyy committed
534
535
    if optimizer is not None and isinstance(optimizer, DummyOptimizer):
        optimizer_dict = optimizer.optimizer_dict
chenych's avatar
chenych committed
536
        scheduler_dict: dict[torch.nn.Parameter, torch.optim.lr_scheduler.LRScheduler] = {}
Rayyyyy's avatar
V0.6.3  
Rayyyyy committed
537
538
539
540
541
542
543

        for param in optimizer_dict.keys():
            scheduler_dict[param] = get_scheduler(
                training_args.lr_scheduler_type,
                optimizer=optimizer_dict[param],
                num_warmup_steps=training_args.get_warmup_steps(num_training_steps),
                num_training_steps=num_training_steps,
chenych's avatar
chenych committed
544
                scheduler_specific_kwargs=training_args.lr_scheduler_kwargs,
Rayyyyy's avatar
V0.6.3  
Rayyyyy committed
545
546
547
548
549
550
551
            )

        def scheduler_hook(param: "torch.nn.Parameter"):
            scheduler_dict[param].step()

        for param in optimizer_dict.keys():
            param.register_post_accumulate_grad_hook(scheduler_hook)
chenych's avatar
chenych committed
552
553
554
555


def get_batch_logps(
    logits: "torch.Tensor", labels: "torch.Tensor", label_pad_token_id: int = IGNORE_INDEX
chenych's avatar
chenych committed
556
557
) -> tuple["torch.Tensor", "torch.Tensor"]:
    r"""Compute the log probabilities of the given labels under the given logits.
chenych's avatar
chenych committed
558
559
560
561

    Returns:
        logps: A tensor of shape (batch_size,) containing the sum of log probabilities.
        valid_length: A tensor of shape (batch_size,) containing the number of non-masked tokens.
chenych's avatar
chenych committed
562

chenych's avatar
chenych committed
563
564
565
566
567
568
569
570
571
572
    """
    if logits.shape[:-1] != labels.shape:
        raise ValueError("Logits (batchsize x seqlen) and labels must have the same shape.")

    labels = labels[:, 1:].clone()
    logits = logits[:, :-1, :]
    loss_mask = labels != label_pad_token_id
    labels[labels == label_pad_token_id] = 0  # dummy token
    per_token_logps = torch.gather(logits.log_softmax(-1), dim=2, index=labels.unsqueeze(2)).squeeze(2)
    return (per_token_logps * loss_mask).sum(-1), loss_mask.sum(-1)
luopl's avatar
luopl committed
573
574
575


def nested_detach(
chenych's avatar
chenych committed
576
    tensors: Union["torch.Tensor", list["torch.Tensor"], tuple["torch.Tensor"], dict[str, "torch.Tensor"]],
luopl's avatar
luopl committed
577
578
    clone: bool = False,
):
chenych's avatar
chenych committed
579
    r"""Detach `tensors` (even if it's a nested list/tuple/dict of tensors)."""
luopl's avatar
luopl committed
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
    if isinstance(tensors, (list, tuple)):
        return type(tensors)(nested_detach(t, clone=clone) for t in tensors)
    elif isinstance(tensors, Mapping):
        return type(tensors)({k: nested_detach(t, clone=clone) for k, t in tensors.items()})

    if isinstance(tensors, torch.Tensor):
        if clone:
            return tensors.detach().clone()
        else:
            return tensors.detach()
    else:
        return tensors


def get_swanlab_callback(finetuning_args: "FinetuningArguments") -> "TrainerCallback":
chenych's avatar
chenych committed
595
    r"""Get the callback for logging to SwanLab."""
luopl's avatar
luopl committed
596
597
598
599
600
601
    import swanlab  # type: ignore
    from swanlab.integration.transformers import SwanLabCallback  # type: ignore

    if finetuning_args.swanlab_api_key is not None:
        swanlab.login(api_key=finetuning_args.swanlab_api_key)

chenych's avatar
chenych committed
602
603
604
605
606
607
608
609
610
    if finetuning_args.swanlab_lark_webhook_url is not None:
        from swanlab.plugin.notification import LarkCallback  # type: ignore

        lark_callback = LarkCallback(
            webhook_url=finetuning_args.swanlab_lark_webhook_url,
            secret=finetuning_args.swanlab_lark_secret,
        )
        swanlab.register_callbacks([lark_callback])

chenych's avatar
chenych committed
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
    class SwanLabCallbackExtension(SwanLabCallback):
        def setup(self, args: "TrainingArguments", state: "TrainerState", model: "PreTrainedModel", **kwargs):
            if not state.is_world_process_zero:
                return

            super().setup(args, state, model, **kwargs)
            try:
                if hasattr(self, "_swanlab"):
                    swanlab_public_config = self._swanlab.get_run().public.json()
                else:  # swanlab <= 0.4.9
                    swanlab_public_config = self._experiment.get_run().public.json()
            except Exception:
                swanlab_public_config = {}

            with open(os.path.join(args.output_dir, SWANLAB_CONFIG), "w") as f:
                f.write(json.dumps(swanlab_public_config, indent=2))

    swanlab_callback = SwanLabCallbackExtension(
luopl's avatar
luopl committed
629
630
631
632
633
        project=finetuning_args.swanlab_project,
        workspace=finetuning_args.swanlab_workspace,
        experiment_name=finetuning_args.swanlab_run_name,
        mode=finetuning_args.swanlab_mode,
        config={"Framework": "🦙LlamaFactory"},
chenych's avatar
chenych committed
634
        logdir=finetuning_args.swanlab_logdir,
luopl's avatar
luopl committed
635
636
637
638
639
640
    )
    return swanlab_callback


def get_ray_trainer(
    training_function: Callable,
chenych's avatar
chenych committed
641
    train_loop_config: dict[str, Any],
luopl's avatar
luopl committed
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
    ray_args: "RayArguments",
) -> "TorchTrainer":
    if not ray_args.use_ray:
        raise ValueError("Ray was not enabled. Please set `USE_RAY=1` to enable ray.")

    trainer = TorchTrainer(
        training_function,
        train_loop_config=train_loop_config,
        scaling_config=ScalingConfig(
            num_workers=ray_args.ray_num_workers,
            resources_per_worker=ray_args.resources_per_worker,
            placement_strategy=ray_args.placement_strategy,
            use_gpu=True,
        ),
        run_config=RunConfig(
            name=ray_args.ray_run_name,
chenych's avatar
chenych committed
658
            storage_path=Path(ray_args.ray_storage_path).absolute().as_posix(),
luopl's avatar
luopl committed
659
660
661
        ),
    )
    return trainer