__init__.py 2.1 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061
  1. from typing import Any
  2. def build_lora_config(params: dict[str, Any]):
  3. """返回实际的 peft.LoraConfig 对象。"""
  4. from peft import LoraConfig, TaskType
  5. target_modules = params.get("lora_target_modules", "all-linear")
  6. if isinstance(target_modules, str):
  7. if target_modules == "all-linear":
  8. target_modules = ["linear", "lm_head", "q_proj", "v_proj", "k_proj", "o_proj"]
  9. return LoraConfig(
  10. r=params.get("lora_r", 16),
  11. lora_alpha=params.get("lora_alpha", 32),
  12. lora_dropout=params.get("lora_dropout", 0.05),
  13. target_modules=target_modules,
  14. task_type=TaskType.CAUSAL_LM,
  15. )
  16. def build_qlora_config(params: dict[str, Any]):
  17. """返回 (bitsandbytes_config, peft.LoraConfig) 二元组。"""
  18. from peft import LoraConfig, TaskType
  19. from transformers import BitsAndBytesConfig
  20. import torch
  21. bnb_params = BitsAndBytesConfig(
  22. load_in_4bit=params.get("qlora_bits", 4) == 4,
  23. load_in_8bit=params.get("qlora_bits", 4) == 8,
  24. bnb_4bit_quant_type=params.get("qlora_type", "nf4"),
  25. bnb_4bit_use_double_quant=params.get("qlora_double_quant", True),
  26. bnb_4bit_compute_dtype=torch.float16,
  27. )
  28. target_modules = params.get("lora_target_modules", "all-linear")
  29. if isinstance(target_modules, str) and target_modules == "all-linear":
  30. target_modules = ["linear", "lm_head", "q_proj", "v_proj", "k_proj", "o_proj"]
  31. lora_cfg = LoraConfig(
  32. r=params.get("lora_r", 16),
  33. lora_alpha=params.get("lora_alpha", 32),
  34. lora_dropout=params.get("lora_dropout", 0.05),
  35. target_modules=target_modules,
  36. task_type=TaskType.CAUSAL_LM,
  37. )
  38. return bnb_params, lora_cfg
  39. def build_adalora_config(params: dict[str, Any]):
  40. """返回实际的 peft.AdaLoraConfig 对象。"""
  41. from peft import AdaLoraConfig, TaskType
  42. return AdaLoraConfig(
  43. init_r=params.get("adalora_init_r", 8),
  44. target_r=params.get("adalora_target_r", 16),
  45. beta1=params.get("adalora_beta1", 0.85),
  46. beta2=params.get("adalora_beta2", 0.85),
  47. task_type=TaskType.CAUSAL_LM,
  48. )