parser.py 20.5 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
#
# This code is inspired by the HuggingFace's transformers library.
# https://github.com/huggingface/transformers/blob/v4.40.0/examples/pytorch/language-modeling/run_clm.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.

luopl's avatar
luopl committed
18
import json
Rayyyyy's avatar
V0.6.3  
Rayyyyy committed
19
20
import os
import sys
luopl's avatar
luopl committed
21
from pathlib import Path
chenych's avatar
chenych committed
22
from typing import Any, Optional, Union
Rayyyyy's avatar
V0.6.3  
Rayyyyy committed
23
24
25

import torch
import transformers
luopl's avatar
luopl committed
26
import yaml
chenych's avatar
chenych committed
27
from omegaconf import OmegaConf
luopl's avatar
luopl committed
28
from transformers import HfArgumentParser
chenych's avatar
chenych committed
29
from transformers.integrations import is_deepspeed_zero3_enabled
Rayyyyy's avatar
V0.6.3  
Rayyyyy committed
30
from transformers.trainer_utils import get_last_checkpoint
chenych's avatar
chenych committed
31
from transformers.training_args import ParallelMode
luopl's avatar
luopl committed
32
from transformers.utils import is_torch_bf16_gpu_available, is_torch_npu_available
Rayyyyy's avatar
V0.6.3  
Rayyyyy committed
33

luopl's avatar
luopl committed
34
from ..extras import logging
chenych's avatar
chenych committed
35
from ..extras.constants import CHECKPOINT_NAMES, EngineName
chenych's avatar
chenych committed
36
from ..extras.misc import check_dependencies, check_version, get_current_device, is_env_enabled
Rayyyyy's avatar
V0.6.3  
Rayyyyy committed
37
38
39
40
41
from .data_args import DataArguments
from .evaluation_args import EvaluationArguments
from .finetuning_args import FinetuningArguments
from .generating_args import GeneratingArguments
from .model_args import ModelArguments
luopl's avatar
luopl committed
42
from .training_args import RayArguments, TrainingArguments
Rayyyyy's avatar
V0.6.3  
Rayyyyy committed
43
44


luopl's avatar
luopl committed
45
logger = logging.get_logger(__name__)
Rayyyyy's avatar
V0.6.3  
Rayyyyy committed
46
47
48
49

check_dependencies()


luopl's avatar
luopl committed
50
_TRAIN_ARGS = [ModelArguments, DataArguments, TrainingArguments, FinetuningArguments, GeneratingArguments]
chenych's avatar
chenych committed
51
_TRAIN_CLS = tuple[ModelArguments, DataArguments, TrainingArguments, FinetuningArguments, GeneratingArguments]
Rayyyyy's avatar
V0.6.3  
Rayyyyy committed
52
_INFER_ARGS = [ModelArguments, DataArguments, FinetuningArguments, GeneratingArguments]
chenych's avatar
chenych committed
53
_INFER_CLS = tuple[ModelArguments, DataArguments, FinetuningArguments, GeneratingArguments]
Rayyyyy's avatar
V0.6.3  
Rayyyyy committed
54
_EVAL_ARGS = [ModelArguments, DataArguments, EvaluationArguments, FinetuningArguments]
chenych's avatar
chenych committed
55
_EVAL_CLS = tuple[ModelArguments, DataArguments, EvaluationArguments, FinetuningArguments]
Rayyyyy's avatar
V0.6.3  
Rayyyyy committed
56
57


chenych's avatar
chenych committed
58
59
def read_args(args: Optional[Union[dict[str, Any], list[str]]] = None) -> Union[dict[str, Any], list[str]]:
    r"""Get arguments from the command line or a config file."""
Rayyyyy's avatar
V0.6.3  
Rayyyyy committed
60
    if args is not None:
luopl's avatar
luopl committed
61
        return args
Rayyyyy's avatar
V0.6.3  
Rayyyyy committed
62

chenych's avatar
chenych committed
63
64
65
66
67
68
69
70
    if sys.argv[1].endswith(".yaml") or sys.argv[1].endswith(".yml"):
        override_config = OmegaConf.from_cli(sys.argv[2:])
        dict_config = yaml.safe_load(Path(sys.argv[1]).absolute().read_text())
        return OmegaConf.to_container(OmegaConf.merge(dict_config, override_config))
    elif sys.argv[1].endswith(".json"):
        override_config = OmegaConf.from_cli(sys.argv[2:])
        dict_config = json.loads(Path(sys.argv[1]).absolute().read_text())
        return OmegaConf.to_container(OmegaConf.merge(dict_config, override_config))
luopl's avatar
luopl committed
71
72
    else:
        return sys.argv[1:]
Rayyyyy's avatar
V0.6.3  
Rayyyyy committed
73
74


luopl's avatar
luopl committed
75
def _parse_args(
chenych's avatar
chenych committed
76
77
    parser: "HfArgumentParser", args: Optional[Union[dict[str, Any], list[str]]] = None, allow_extra_keys: bool = False
) -> tuple[Any]:
luopl's avatar
luopl committed
78
79
80
    args = read_args(args)
    if isinstance(args, dict):
        return parser.parse_dict(args, allow_extra_keys=allow_extra_keys)
Rayyyyy's avatar
V0.6.3  
Rayyyyy committed
81

luopl's avatar
luopl committed
82
83
84
    (*parsed_args, unknown_args) = parser.parse_args_into_dataclasses(args=args, return_remaining_strings=True)

    if unknown_args and not allow_extra_keys:
Rayyyyy's avatar
V0.6.3  
Rayyyyy committed
85
        print(parser.format_help())
luopl's avatar
luopl committed
86
87
        print(f"Got unknown args, potentially deprecated arguments: {unknown_args}")
        raise ValueError(f"Some specified arguments are not used by the HfArgumentParser: {unknown_args}")
Rayyyyy's avatar
V0.6.3  
Rayyyyy committed
88

chenych's avatar
chenych committed
89
    return tuple(parsed_args)
Rayyyyy's avatar
V0.6.3  
Rayyyyy committed
90
91


luopl's avatar
luopl committed
92
def _set_transformers_logging() -> None:
chenych's avatar
chenych committed
93
94
95
96
    if os.getenv("LLAMAFACTORY_VERBOSITY", "INFO") in ["DEBUG", "INFO"]:
        transformers.utils.logging.set_verbosity_info()
        transformers.utils.logging.enable_default_handler()
        transformers.utils.logging.enable_explicit_format()
Rayyyyy's avatar
V0.6.3  
Rayyyyy committed
97
98


chenych's avatar
chenych committed
99
100
101
102
103
104
105
106
def _set_env_vars() -> None:
    if is_torch_npu_available():
        # avoid JIT compile on NPU devices, see https://zhuanlan.zhihu.com/p/660875458
        torch.npu.set_compile_mode(jit_compile=is_env_enabled("NPU_JIT_COMPILE"))
        # avoid use fork method on NPU devices, see https://github.com/hiyouga/LLaMA-Factory/issues/7447
        os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"


chenych's avatar
chenych committed
107
108
109
110
111
def _verify_model_args(
    model_args: "ModelArguments",
    data_args: "DataArguments",
    finetuning_args: "FinetuningArguments",
) -> None:
Rayyyyy's avatar
V0.6.3  
Rayyyyy committed
112
113
114
115
116
117
118
    if model_args.adapter_name_or_path is not None and finetuning_args.finetuning_type != "lora":
        raise ValueError("Adapter is only valid for the LoRA method.")

    if model_args.quantization_bit is not None:
        if finetuning_args.finetuning_type != "lora":
            raise ValueError("Quantization is only compatible with the LoRA method.")

chenych's avatar
chenych committed
119
120
121
122
123
124
        if finetuning_args.pissa_init:
            raise ValueError("Please use scripts/pissa_init.py to initialize PiSSA for a quantized model.")

        if model_args.resize_vocab:
            raise ValueError("Cannot resize embedding layers of a quantized model.")

Rayyyyy's avatar
V0.6.3  
Rayyyyy committed
125
126
127
128
129
130
        if model_args.adapter_name_or_path is not None and finetuning_args.create_new_adapter:
            raise ValueError("Cannot create new adapter upon a quantized model.")

        if model_args.adapter_name_or_path is not None and len(model_args.adapter_name_or_path) != 1:
            raise ValueError("Quantized model only accepts a single adapter. Merge them first.")

chenych's avatar
chenych committed
131
    if data_args.template == "yi" and model_args.use_fast_tokenizer:
luopl's avatar
luopl committed
132
        logger.warning_rank0("We should use slow tokenizer for the Yi models. Change `use_fast_tokenizer` to False.")
chenych's avatar
chenych committed
133
134
        model_args.use_fast_tokenizer = False

Rayyyyy's avatar
V0.6.3  
Rayyyyy committed
135
136
137
138

def _check_extra_dependencies(
    model_args: "ModelArguments",
    finetuning_args: "FinetuningArguments",
luopl's avatar
luopl committed
139
    training_args: Optional["TrainingArguments"] = None,
Rayyyyy's avatar
V0.6.3  
Rayyyyy committed
140
141
) -> None:
    if model_args.use_unsloth:
luopl's avatar
luopl committed
142
        check_version("unsloth", mandatory=True)
Rayyyyy's avatar
V0.6.3  
Rayyyyy committed
143

luopl's avatar
luopl committed
144
    if model_args.enable_liger_kernel:
luopl's avatar
luopl committed
145
        check_version("liger-kernel", mandatory=True)
luopl's avatar
luopl committed
146

Rayyyyy's avatar
V0.6.3  
Rayyyyy committed
147
    if model_args.mixture_of_depths is not None:
luopl's avatar
luopl committed
148
        check_version("mixture-of-depth>=1.1.6", mandatory=True)
Rayyyyy's avatar
V0.6.3  
Rayyyyy committed
149

chenych's avatar
chenych committed
150
    if model_args.infer_backend == EngineName.VLLM:
chenych's avatar
chenych committed
151
        check_version("vllm>=0.4.3,<=0.9.1")
luopl's avatar
luopl committed
152
        check_version("vllm", mandatory=True)
chenych's avatar
chenych committed
153
    elif model_args.infer_backend == EngineName.SGLANG:
chenych's avatar
chenych committed
154
        check_version("sglang>=0.4.5")
chenych's avatar
chenych committed
155
        check_version("sglang", mandatory=True)
Rayyyyy's avatar
V0.6.3  
Rayyyyy committed
156
157

    if finetuning_args.use_galore:
luopl's avatar
luopl committed
158
159
160
161
        check_version("galore_torch", mandatory=True)

    if finetuning_args.use_apollo:
        check_version("apollo_torch", mandatory=True)
Rayyyyy's avatar
V0.6.3  
Rayyyyy committed
162
163

    if finetuning_args.use_badam:
luopl's avatar
luopl committed
164
        check_version("badam>=1.2.1", mandatory=True)
chenych's avatar
chenych committed
165
166

    if finetuning_args.use_adam_mini:
luopl's avatar
luopl committed
167
        check_version("adam-mini", mandatory=True)
chenych's avatar
chenych committed
168
169

    if finetuning_args.plot_loss:
luopl's avatar
luopl committed
170
        check_version("matplotlib", mandatory=True)
Rayyyyy's avatar
V0.6.3  
Rayyyyy committed
171

chenych's avatar
chenych committed
172
173
174
175
176
177
178
179
180
    if training_args is not None:
        if training_args.deepspeed:
            # pin deepspeed version < 0.17 because of https://github.com/deepspeedai/DeepSpeed/issues/7347
            check_version("deepspeed>=0.10.0,<=0.16.9", mandatory=True)

        if training_args.predict_with_generate:
            check_version("jieba", mandatory=True)
            check_version("nltk", mandatory=True)
            check_version("rouge_chinese", mandatory=True)
Rayyyyy's avatar
V0.6.3  
Rayyyyy committed
181
182


chenych's avatar
chenych committed
183
def _parse_train_args(args: Optional[Union[dict[str, Any], list[str]]] = None) -> _TRAIN_CLS:
Rayyyyy's avatar
V0.6.3  
Rayyyyy committed
184
    parser = HfArgumentParser(_TRAIN_ARGS)
chenych's avatar
chenych committed
185
186
    allow_extra_keys = is_env_enabled("ALLOW_EXTRA_ARGS")
    return _parse_args(parser, args, allow_extra_keys=allow_extra_keys)
Rayyyyy's avatar
V0.6.3  
Rayyyyy committed
187
188


chenych's avatar
chenych committed
189
def _parse_infer_args(args: Optional[Union[dict[str, Any], list[str]]] = None) -> _INFER_CLS:
Rayyyyy's avatar
V0.6.3  
Rayyyyy committed
190
    parser = HfArgumentParser(_INFER_ARGS)
chenych's avatar
chenych committed
191
192
    allow_extra_keys = is_env_enabled("ALLOW_EXTRA_ARGS")
    return _parse_args(parser, args, allow_extra_keys=allow_extra_keys)
Rayyyyy's avatar
V0.6.3  
Rayyyyy committed
193
194


chenych's avatar
chenych committed
195
def _parse_eval_args(args: Optional[Union[dict[str, Any], list[str]]] = None) -> _EVAL_CLS:
Rayyyyy's avatar
V0.6.3  
Rayyyyy committed
196
    parser = HfArgumentParser(_EVAL_ARGS)
chenych's avatar
chenych committed
197
198
    allow_extra_keys = is_env_enabled("ALLOW_EXTRA_ARGS")
    return _parse_args(parser, args, allow_extra_keys=allow_extra_keys)
Rayyyyy's avatar
V0.6.3  
Rayyyyy committed
199
200


chenych's avatar
chenych committed
201
def get_ray_args(args: Optional[Union[dict[str, Any], list[str]]] = None) -> RayArguments:
luopl's avatar
luopl committed
202
203
204
205
206
    parser = HfArgumentParser(RayArguments)
    (ray_args,) = _parse_args(parser, args, allow_extra_keys=True)
    return ray_args


chenych's avatar
chenych committed
207
def get_train_args(args: Optional[Union[dict[str, Any], list[str]]] = None) -> _TRAIN_CLS:
Rayyyyy's avatar
V0.6.3  
Rayyyyy committed
208
209
210
211
212
213
214
    model_args, data_args, training_args, finetuning_args, generating_args = _parse_train_args(args)

    # Setup logging
    if training_args.should_log:
        _set_transformers_logging()

    # Check arguments
chenych's avatar
chenych committed
215
216
217
218
219
220
221
222
223
    if finetuning_args.stage != "sft":
        if training_args.predict_with_generate:
            raise ValueError("`predict_with_generate` cannot be set as True except SFT.")

        if data_args.neat_packing:
            raise ValueError("`neat_packing` cannot be set as True except SFT.")

        if data_args.train_on_prompt or data_args.mask_history:
            raise ValueError("`train_on_prompt` or `mask_history` cannot be set as True except SFT.")
Rayyyyy's avatar
V0.6.3  
Rayyyyy committed
224
225
226
227
228
229
230

    if finetuning_args.stage == "sft" and training_args.do_predict and not training_args.predict_with_generate:
        raise ValueError("Please enable `predict_with_generate` to save model predictions.")

    if finetuning_args.stage in ["rm", "ppo"] and training_args.load_best_model_at_end:
        raise ValueError("RM and PPO stages do not support `load_best_model_at_end`.")

chenych's avatar
chenych committed
231
232
233
    if finetuning_args.stage == "ppo":
        if not training_args.do_train:
            raise ValueError("PPO training does not support evaluation, use the SFT stage to evaluate models.")
Rayyyyy's avatar
V0.6.3  
Rayyyyy committed
234

chenych's avatar
chenych committed
235
236
        if model_args.shift_attn:
            raise ValueError("PPO training is incompatible with S^2-Attn.")
Rayyyyy's avatar
V0.6.3  
Rayyyyy committed
237

chenych's avatar
chenych committed
238
239
        if finetuning_args.reward_model_type == "lora" and model_args.use_unsloth:
            raise ValueError("Unsloth does not support lora reward model.")
Rayyyyy's avatar
V0.6.3  
Rayyyyy committed
240

chenych's avatar
chenych committed
241
242
243
244
245
246
247
248
        if training_args.report_to and training_args.report_to[0] not in ["wandb", "tensorboard"]:
            raise ValueError("PPO only accepts wandb or tensorboard logger.")

    if training_args.parallel_mode == ParallelMode.NOT_DISTRIBUTED:
        raise ValueError("Please launch distributed training with `llamafactory-cli` or `torchrun`.")

    if training_args.deepspeed and training_args.parallel_mode != ParallelMode.DISTRIBUTED:
        raise ValueError("Please use `FORCE_TORCHRUN=1` to launch DeepSpeed training.")
Rayyyyy's avatar
V0.6.3  
Rayyyyy committed
249
250
251
252

    if training_args.max_steps == -1 and data_args.streaming:
        raise ValueError("Please specify `max_steps` in streaming mode.")

chenych's avatar
chenych committed
253
254
255
256
257
258
259
260
    if training_args.do_train and data_args.dataset is None:
        raise ValueError("Please specify dataset for training.")

    if (training_args.do_eval or training_args.do_predict) and (
        data_args.eval_dataset is None and data_args.val_size < 1e-6
    ):
        raise ValueError("Please specify dataset for evaluation.")

luopl's avatar
luopl committed
261
262
263
264
265
266
    if training_args.predict_with_generate:
        if is_deepspeed_zero3_enabled():
            raise ValueError("`predict_with_generate` is incompatible with DeepSpeed ZeRO-3.")

        if data_args.eval_dataset is None:
            raise ValueError("Cannot use `predict_with_generate` if `eval_dataset` is None.")
chenych's avatar
chenych committed
267

luopl's avatar
luopl committed
268
269
        if finetuning_args.compute_accuracy:
            raise ValueError("Cannot use `predict_with_generate` and `compute_accuracy` together.")
Rayyyyy's avatar
V0.6.3  
Rayyyyy committed
270
271
272
273

    if training_args.do_train and model_args.quantization_device_map == "auto":
        raise ValueError("Cannot use device map for quantized models in training.")

chenych's avatar
chenych committed
274
275
    if finetuning_args.pissa_init and is_deepspeed_zero3_enabled():
        raise ValueError("Please use scripts/pissa_init.py to initialize PiSSA in DeepSpeed ZeRO-3.")
Rayyyyy's avatar
V0.6.3  
Rayyyyy committed
276
277

    if finetuning_args.pure_bf16:
luopl's avatar
luopl committed
278
        if not (is_torch_bf16_gpu_available() or (is_torch_npu_available() and torch.npu.is_bf16_supported())):
Rayyyyy's avatar
V0.6.3  
Rayyyyy committed
279
280
            raise ValueError("This device does not support `pure_bf16`.")

chenych's avatar
chenych committed
281
282
        if is_deepspeed_zero3_enabled():
            raise ValueError("`pure_bf16` is incompatible with DeepSpeed ZeRO-3.")
Rayyyyy's avatar
V0.6.3  
Rayyyyy committed
283

luopl's avatar
luopl committed
284
285
286
287
288
289
    if training_args.parallel_mode == ParallelMode.DISTRIBUTED:
        if finetuning_args.use_galore and finetuning_args.galore_layerwise:
            raise ValueError("Distributed training does not support layer-wise GaLore.")

        if finetuning_args.use_apollo and finetuning_args.apollo_layerwise:
            raise ValueError("Distributed training does not support layer-wise APOLLO.")
Rayyyyy's avatar
V0.6.3  
Rayyyyy committed
290

luopl's avatar
luopl committed
291
292
293
294
295
        if finetuning_args.use_badam:
            if finetuning_args.badam_mode == "ratio":
                raise ValueError("Radio-based BAdam does not yet support distributed training, use layer-wise BAdam.")
            elif not is_deepspeed_zero3_enabled():
                raise ValueError("Layer-wise BAdam only supports DeepSpeed ZeRO-3 training.")
Rayyyyy's avatar
V0.6.3  
Rayyyyy committed
296

luopl's avatar
luopl committed
297
298
    if training_args.deepspeed is not None and (finetuning_args.use_galore or finetuning_args.use_apollo):
        raise ValueError("GaLore and APOLLO are incompatible with DeepSpeed yet.")
Rayyyyy's avatar
V0.6.3  
Rayyyyy committed
299

chenych's avatar
chenych committed
300
301
    if model_args.infer_backend != EngineName.HF:
        raise ValueError("vLLM/SGLang backend is only available for API, CLI and Web.")
Rayyyyy's avatar
V0.6.3  
Rayyyyy committed
302

chenych's avatar
chenych committed
303
304
305
    if model_args.use_unsloth and is_deepspeed_zero3_enabled():
        raise ValueError("Unsloth is incompatible with DeepSpeed ZeRO-3.")

chenych's avatar
chenych committed
306
    _set_env_vars()
chenych's avatar
chenych committed
307
    _verify_model_args(model_args, data_args, finetuning_args)
Rayyyyy's avatar
V0.6.3  
Rayyyyy committed
308
309
310
311
312
    _check_extra_dependencies(model_args, finetuning_args, training_args)

    if (
        training_args.do_train
        and finetuning_args.finetuning_type == "lora"
chenych's avatar
chenych committed
313
        and model_args.quantization_bit is None
Rayyyyy's avatar
V0.6.3  
Rayyyyy committed
314
315
316
        and model_args.resize_vocab
        and finetuning_args.additional_target is None
    ):
luopl's avatar
luopl committed
317
318
319
        logger.warning_rank0(
            "Remember to add embedding layers to `additional_target` to make the added tokens trainable."
        )
Rayyyyy's avatar
V0.6.3  
Rayyyyy committed
320
321

    if training_args.do_train and model_args.quantization_bit is not None and (not model_args.upcast_layernorm):
luopl's avatar
luopl committed
322
        logger.warning_rank0("We recommend enable `upcast_layernorm` in quantized training.")
Rayyyyy's avatar
V0.6.3  
Rayyyyy committed
323
324

    if training_args.do_train and (not training_args.fp16) and (not training_args.bf16):
luopl's avatar
luopl committed
325
        logger.warning_rank0("We recommend enable mixed precision training.")
Rayyyyy's avatar
V0.6.3  
Rayyyyy committed
326

luopl's avatar
luopl committed
327
328
329
330
331
    if (
        training_args.do_train
        and (finetuning_args.use_galore or finetuning_args.use_apollo)
        and not finetuning_args.pure_bf16
    ):
luopl's avatar
luopl committed
332
        logger.warning_rank0(
luopl's avatar
luopl committed
333
            "Using GaLore or APOLLO with mixed precision training may significantly increases GPU memory usage."
luopl's avatar
luopl committed
334
        )
Rayyyyy's avatar
V0.6.3  
Rayyyyy committed
335
336

    if (not training_args.do_train) and model_args.quantization_bit is not None:
luopl's avatar
luopl committed
337
        logger.warning_rank0("Evaluating model in 4/8-bit mode may cause lower scores.")
Rayyyyy's avatar
V0.6.3  
Rayyyyy committed
338
339

    if (not training_args.do_train) and finetuning_args.stage == "dpo" and finetuning_args.ref_model is None:
luopl's avatar
luopl committed
340
        logger.warning_rank0("Specify `ref_model` for computing rewards at evaluation.")
Rayyyyy's avatar
V0.6.3  
Rayyyyy committed
341
342

    # Post-process training arguments
chenych's avatar
chenych committed
343
344
345
346
347
348
349
350
    training_args.generation_max_length = training_args.generation_max_length or data_args.cutoff_len
    training_args.generation_num_beams = data_args.eval_num_beams or training_args.generation_num_beams
    training_args.remove_unused_columns = False  # important for multimodal dataset

    if finetuning_args.finetuning_type == "lora":
        # https://github.com/huggingface/transformers/blob/v4.50.0/src/transformers/trainer.py#L782
        training_args.label_names = training_args.label_names or ["labels"]

Rayyyyy's avatar
V0.6.3  
Rayyyyy committed
351
    if (
chenych's avatar
chenych committed
352
        training_args.parallel_mode == ParallelMode.DISTRIBUTED
Rayyyyy's avatar
V0.6.3  
Rayyyyy committed
353
354
355
        and training_args.ddp_find_unused_parameters is None
        and finetuning_args.finetuning_type == "lora"
    ):
chenych's avatar
chenych committed
356
        logger.info_rank0("Set `ddp_find_unused_parameters` to False in DDP training since LoRA is enabled.")
Rayyyyy's avatar
V0.6.3  
Rayyyyy committed
357
358
359
360
361
        training_args.ddp_find_unused_parameters = False

    if finetuning_args.stage in ["rm", "ppo"] and finetuning_args.finetuning_type in ["full", "freeze"]:
        can_resume_from_checkpoint = False
        if training_args.resume_from_checkpoint is not None:
luopl's avatar
luopl committed
362
            logger.warning_rank0("Cannot resume from checkpoint in current stage.")
Rayyyyy's avatar
V0.6.3  
Rayyyyy committed
363
364
365
366
367
368
369
370
371
372
373
374
            training_args.resume_from_checkpoint = None
    else:
        can_resume_from_checkpoint = True

    if (
        training_args.resume_from_checkpoint is None
        and training_args.do_train
        and os.path.isdir(training_args.output_dir)
        and not training_args.overwrite_output_dir
        and can_resume_from_checkpoint
    ):
        last_checkpoint = get_last_checkpoint(training_args.output_dir)
chenych's avatar
chenych committed
375
376
377
        if last_checkpoint is None and any(
            os.path.isfile(os.path.join(training_args.output_dir, name)) for name in CHECKPOINT_NAMES
        ):
Rayyyyy's avatar
V0.6.3  
Rayyyyy committed
378
379
380
381
            raise ValueError("Output directory already exists and is not empty. Please set `overwrite_output_dir`.")

        if last_checkpoint is not None:
            training_args.resume_from_checkpoint = last_checkpoint
luopl's avatar
luopl committed
382
383
            logger.info_rank0(f"Resuming training from {training_args.resume_from_checkpoint}.")
            logger.info_rank0("Change `output_dir` or use `overwrite_output_dir` to avoid.")
Rayyyyy's avatar
V0.6.3  
Rayyyyy committed
384
385
386
387
388
389

    if (
        finetuning_args.stage in ["rm", "ppo"]
        and finetuning_args.finetuning_type == "lora"
        and training_args.resume_from_checkpoint is not None
    ):
luopl's avatar
luopl committed
390
        logger.warning_rank0(
chenych's avatar
chenych committed
391
            f"Add {training_args.resume_from_checkpoint} to `adapter_name_or_path` to resume training from checkpoint."
Rayyyyy's avatar
V0.6.3  
Rayyyyy committed
392
393
394
395
396
397
398
399
400
401
        )

    # Post-process model arguments
    if training_args.bf16 or finetuning_args.pure_bf16:
        model_args.compute_dtype = torch.bfloat16
    elif training_args.fp16:
        model_args.compute_dtype = torch.float16

    model_args.device_map = {"": get_current_device()}
    model_args.model_max_length = data_args.cutoff_len
chenych's avatar
chenych committed
402
    model_args.block_diag_attn = data_args.neat_packing
Rayyyyy's avatar
V0.6.3  
Rayyyyy committed
403
404
    data_args.packing = data_args.packing if data_args.packing is not None else finetuning_args.stage == "pt"

chenych's avatar
chenych committed
405
    # Log on each process the small summary
Rayyyyy's avatar
V0.6.3  
Rayyyyy committed
406
    logger.info(
chenych's avatar
chenych committed
407
408
409
410
        f"Process rank: {training_args.process_index}, "
        f"world size: {training_args.world_size}, device: {training_args.device}, "
        f"distributed training: {training_args.parallel_mode == ParallelMode.DISTRIBUTED}, "
        f"compute dtype: {str(model_args.compute_dtype)}"
Rayyyyy's avatar
V0.6.3  
Rayyyyy committed
411
412
413
414
415
416
    )
    transformers.set_seed(training_args.seed)

    return model_args, data_args, training_args, finetuning_args, generating_args


chenych's avatar
chenych committed
417
def get_infer_args(args: Optional[Union[dict[str, Any], list[str]]] = None) -> _INFER_CLS:
Rayyyyy's avatar
V0.6.3  
Rayyyyy committed
418
419
    model_args, data_args, finetuning_args, generating_args = _parse_infer_args(args)

chenych's avatar
chenych committed
420
    # Setup logging
Rayyyyy's avatar
V0.6.3  
Rayyyyy committed
421
422
    _set_transformers_logging()

chenych's avatar
chenych committed
423
    # Check arguments
Rayyyyy's avatar
V0.6.3  
Rayyyyy committed
424
425
426
427
428
    if model_args.infer_backend == "vllm":
        if finetuning_args.stage != "sft":
            raise ValueError("vLLM engine only supports auto-regressive models.")

        if model_args.quantization_bit is not None:
chenych's avatar
chenych committed
429
            raise ValueError("vLLM engine does not support bnb quantization (GPTQ and AWQ are supported).")
Rayyyyy's avatar
V0.6.3  
Rayyyyy committed
430
431
432
433

        if model_args.rope_scaling is not None:
            raise ValueError("vLLM engine does not support RoPE scaling.")

chenych's avatar
chenych committed
434
435
436
        if model_args.adapter_name_or_path is not None and len(model_args.adapter_name_or_path) != 1:
            raise ValueError("vLLM only accepts a single adapter. Merge them first.")

chenych's avatar
chenych committed
437
    _set_env_vars()
chenych's avatar
chenych committed
438
    _verify_model_args(model_args, data_args, finetuning_args)
Rayyyyy's avatar
V0.6.3  
Rayyyyy committed
439
440
    _check_extra_dependencies(model_args, finetuning_args)

chenych's avatar
chenych committed
441
    # Post-process model arguments
chenych's avatar
chenych committed
442
443
    if model_args.export_dir is not None and model_args.export_device == "cpu":
        model_args.device_map = {"": torch.device("cpu")}
chenych's avatar
chenych committed
444
445
        if data_args.cutoff_len != DataArguments().cutoff_len:  # override cutoff_len if it is not default
            model_args.model_max_length = data_args.cutoff_len
Rayyyyy's avatar
V0.6.3  
Rayyyyy committed
446
447
448
449
450
451
    else:
        model_args.device_map = "auto"

    return model_args, data_args, finetuning_args, generating_args


chenych's avatar
chenych committed
452
def get_eval_args(args: Optional[Union[dict[str, Any], list[str]]] = None) -> _EVAL_CLS:
Rayyyyy's avatar
V0.6.3  
Rayyyyy committed
453
454
    model_args, data_args, eval_args, finetuning_args = _parse_eval_args(args)

chenych's avatar
chenych committed
455
    # Setup logging
Rayyyyy's avatar
V0.6.3  
Rayyyyy committed
456
457
    _set_transformers_logging()

chenych's avatar
chenych committed
458
    # Check arguments
chenych's avatar
chenych committed
459
460
    if model_args.infer_backend != EngineName.HF:
        raise ValueError("vLLM/SGLang backend is only available for API, CLI and Web.")
Rayyyyy's avatar
V0.6.3  
Rayyyyy committed
461

chenych's avatar
chenych committed
462
    _set_env_vars()
chenych's avatar
chenych committed
463
    _verify_model_args(model_args, data_args, finetuning_args)
Rayyyyy's avatar
V0.6.3  
Rayyyyy committed
464
465
466
467
468
469
470
    _check_extra_dependencies(model_args, finetuning_args)

    model_args.device_map = "auto"

    transformers.set_seed(eval_args.seed)

    return model_args, data_args, eval_args, finetuning_args