support fsdp + qlora

Former-commit-id: b894bf8e84be689db258021f0638e9ac939abcbc
This commit is contained in:
hiyouga
2024-03-21 00:36:06 +08:00
parent 5ed234ca63
commit 935ee0a023
15 changed files with 87 additions and 19 deletions

View File

@@ -65,7 +65,7 @@ def check_dependencies() -> None:
require_version("datasets>=2.14.3", "To fix: pip install datasets>=2.14.3")
require_version("accelerate>=0.27.2", "To fix: pip install accelerate>=0.27.2")
require_version("peft>=0.9.0", "To fix: pip install peft>=0.9.0")
require_version("trl>=0.7.11", "To fix: pip install trl>=0.7.11")
require_version("trl>=0.8.1", "To fix: pip install trl>=0.8.1")
def count_parameters(model: torch.nn.Module) -> Tuple[int, int]:
@@ -81,7 +81,8 @@ def count_parameters(model: torch.nn.Module) -> Tuple[int, int]:
# Due to the design of 4bit linear layers from bitsandbytes, multiply the number of parameters by 2
if param.__class__.__name__ == "Params4bit":
num_params = num_params * 2
num_bytes = param.quant_storage.itemsize if hasattr(param, "quant_storage") else 1
num_params = num_params * 2 * num_bytes
all_param += num_params
if param.requires_grad:

View File

@@ -210,9 +210,6 @@ def _configure_quantization(
logger.info("Quantizing model to {} bit.".format(model_args.export_quantization_bit))
elif model_args.quantization_bit is not None: # bnb
if is_deepspeed_zero3_enabled():
require_version("bitsandbytes>=0.43.0", "To fix: pip install bitsandbytes>=0.43.0")
if model_args.quantization_bit == 8:
require_version("bitsandbytes>=0.37.0", "To fix: pip install bitsandbytes>=0.37.0")
init_kwargs["quantization_config"] = BitsAndBytesConfig(load_in_8bit=True)
@@ -224,6 +221,7 @@ def _configure_quantization(
bnb_4bit_compute_dtype=model_args.compute_dtype,
bnb_4bit_use_double_quant=model_args.double_quantization,
bnb_4bit_quant_type=model_args.quantization_type,
bnb_4bit_quant_storage=model_args.compute_dtype, # crucial for fsdp qlora
)
init_kwargs["device_map"] = {"": get_current_device()}
@@ -300,7 +298,7 @@ def patch_config(
init_kwargs["torch_dtype"] = model_args.compute_dtype
if not is_deepspeed_zero3_enabled():
init_kwargs["low_cpu_mem_usage"] = model_args.low_cpu_mem_usage
if model_args.low_cpu_mem_usage:
if init_kwargs["low_cpu_mem_usage"]:
if "device_map" not in init_kwargs: # quant models cannot use auto device map
init_kwargs["device_map"] = model_args.device_map or {"": get_current_device()}