mirror of
https://github.com/hiyouga/LLaMA-Factory.git
synced 2026-07-28 03:36:10 +08:00
119 lines
4.8 KiB
Python
119 lines
4.8 KiB
Python
# Copyright 2025 HuggingFace Inc., the KVCache.AI team, Approaching AI, and the LlamaFactory team.
|
|
#
|
|
# 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.
|
|
|
|
from dataclasses import dataclass
|
|
from typing import Any, Literal
|
|
|
|
from ...utils import logging
|
|
from ...utils.plugin import BasePlugin
|
|
|
|
|
|
logger = logging.get_logger(__name__)
|
|
|
|
|
|
class QuantizationPlugin(BasePlugin):
|
|
def __call__(
|
|
self,
|
|
init_kwargs: dict[str, Any] | None = None,
|
|
quant_config=None,
|
|
is_trainable: bool = False,
|
|
) -> dict[str, Any]:
|
|
return super().__call__(init_kwargs, quant_config=quant_config, is_trainable=is_trainable)
|
|
|
|
|
|
@dataclass
|
|
class BnbParams:
|
|
name: Literal["bnb", "auto"] = "bnb"
|
|
quantization_bit: int | None = None
|
|
compute_dtype: str | Any = "float16"
|
|
double_quantization: bool = True
|
|
quantization_type: str = "nf4"
|
|
|
|
def __post_init__(self) -> None:
|
|
import torch
|
|
|
|
if isinstance(self.compute_dtype, str):
|
|
dtype = getattr(torch, self.compute_dtype, None)
|
|
if not isinstance(dtype, torch.dtype):
|
|
raise ValueError(f"compute_dtype={self.compute_dtype!r} is not a torch dtype name.")
|
|
self.compute_dtype = dtype
|
|
elif not isinstance(self.compute_dtype, torch.dtype):
|
|
raise TypeError(f"compute_dtype must be str or torch.dtype, got {type(self.compute_dtype).__name__}.")
|
|
|
|
|
|
@QuantizationPlugin("auto").register()
|
|
def quantization_auto(
|
|
init_kwargs: dict[str, Any],
|
|
quant_config: dict | BnbParams,
|
|
is_trainable: bool = False,
|
|
) -> dict[str, Any]:
|
|
quant_config = QuantizationPlugin.parse_params(quant_config, BnbParams)
|
|
if quant_config.quantization_bit is None:
|
|
logger.warning_rank0("No quantization method applied.")
|
|
return init_kwargs
|
|
if quant_config.quantization_bit not in (4, 8):
|
|
raise ValueError(f"Unsupported quantization bit: {quant_config.quantization_bit} for auto quantization.")
|
|
|
|
logger.info_rank0(f"Loading {quant_config.quantization_bit}-bit quantized model.")
|
|
return QuantizationPlugin("bnb")(init_kwargs, quant_config=quant_config, is_trainable=is_trainable)
|
|
|
|
|
|
@QuantizationPlugin("bnb").register()
|
|
def quantization_with_bnb(
|
|
init_kwargs: dict[str, Any],
|
|
quant_config: dict | BnbParams,
|
|
is_trainable: bool = False,
|
|
) -> dict[str, Any]:
|
|
from transformers import BitsAndBytesConfig
|
|
|
|
from ...accelerator.helper import get_current_device
|
|
from ...utils.packages import check_version
|
|
|
|
quant_config = QuantizationPlugin.parse_params(quant_config, BnbParams)
|
|
quantization_bit = quant_config.quantization_bit
|
|
if quantization_bit is None:
|
|
logger.warning_rank0("quantization_bit is not specified, default to 4-bit quantization.")
|
|
quantization_bit = 4
|
|
if quantization_bit not in (4, 8):
|
|
raise ValueError("Bitsandbytes only accepts 4-bit or 8-bit quantization.")
|
|
|
|
logger.info_rank0("Using Bitsandbytes quantization.")
|
|
if quantization_bit == 8:
|
|
check_version("bitsandbytes>=0.37.0", mandatory=True)
|
|
init_kwargs["quantization_config"] = BitsAndBytesConfig(load_in_8bit=True)
|
|
else:
|
|
check_version("bitsandbytes>=0.39.0", mandatory=True)
|
|
init_kwargs["quantization_config"] = BitsAndBytesConfig(
|
|
load_in_4bit=True,
|
|
bnb_4bit_compute_dtype=quant_config.compute_dtype,
|
|
bnb_4bit_use_double_quant=quant_config.double_quantization,
|
|
bnb_4bit_quant_type=quant_config.quantization_type,
|
|
bnb_4bit_quant_storage=quant_config.compute_dtype,
|
|
)
|
|
|
|
if is_trainable:
|
|
logger.info_rank0("Detected inference mode, setting device_map for bitsandbytes quantization.")
|
|
init_kwargs["device_map"] = {"": get_current_device()}
|
|
else:
|
|
logger.info_rank0("Detected training mode, skip setting device_map for bitsandbytes quantization.")
|
|
if quantization_bit != 4:
|
|
raise ValueError("Only 4-bit quantized model can use fsdp+qlora or auto device map.")
|
|
check_version("bitsandbytes>=0.43.0", mandatory=True)
|
|
|
|
logger.info_rank0(f"Quantizing model to {quantization_bit} bit with bitsandbytes.")
|
|
return init_kwargs
|