Load a model from Hugging Face.
(
model_path: str,
device: str,
num_gpus: int,
max_gpu_memory: Optional[str] = None,
load_8bit: bool = False,
cpu_offloading: bool = False,
debug: bool = False,
lora: bool = False,
lora_base_model : str = "huggyllama/llama-7b"
)
| 87 | |
| 88 | |
| 89 | def load_model( |
| 90 | model_path: str, |
| 91 | device: str, |
| 92 | num_gpus: int, |
| 93 | max_gpu_memory: Optional[str] = None, |
| 94 | load_8bit: bool = False, |
| 95 | cpu_offloading: bool = False, |
| 96 | debug: bool = False, |
| 97 | lora: bool = False, |
| 98 | lora_base_model : str = "huggyllama/llama-7b" |
| 99 | ): |
| 100 | """Load a model from Hugging Face.""" |
| 101 | |
| 102 | # Handle device mapping |
| 103 | cpu_offloading = raise_warning_for_incompatible_cpu_offloading_configuration( |
| 104 | device, load_8bit, cpu_offloading |
| 105 | ) |
| 106 | if device == "cpu": |
| 107 | kwargs = {"torch_dtype": torch.float32} |
| 108 | elif device == "cuda": |
| 109 | kwargs = {"torch_dtype": torch.float16} |
| 110 | if lora: |
| 111 | model = LlamaForCausalLM.from_pretrained( |
| 112 | lora_base_model, |
| 113 | load_in_8bit=load_8bit, |
| 114 | torch_dtype=torch.float16, |
| 115 | device_map="auto", |
| 116 | ) |
| 117 | model = PeftModel.from_pretrained( |
| 118 | model, |
| 119 | model_path, |
| 120 | torch_dtype=torch.float16, |
| 121 | ) |
| 122 | |
| 123 | elif num_gpus != 1: |
| 124 | |
| 125 | kwargs["device_map"] = "auto" |
| 126 | if max_gpu_memory is None: |
| 127 | kwargs[ |
| 128 | "device_map" |
| 129 | ] = "sequential" # This is important for not the same VRAM sizes |
| 130 | available_gpu_memory = get_gpu_memory(num_gpus) |
| 131 | kwargs["max_memory"] = { |
| 132 | i: str(int(available_gpu_memory[i] * 0.85)) + "GiB" |
| 133 | for i in range(num_gpus) |
| 134 | } |
| 135 | else: |
| 136 | kwargs["max_memory"] = {i: max_gpu_memory for i in range(num_gpus)} |
| 137 | else: |
| 138 | raise ValueError(f"Invalid device: {device}") |
| 139 | |
| 140 | if cpu_offloading: |
| 141 | # raises an error on incompatible platforms |
| 142 | from transformers import BitsAndBytesConfig |
| 143 | |
| 144 | if "max_memory" in kwargs: |
| 145 | kwargs["max_memory"]["cpu"] = ( |
| 146 | str(math.floor(psutil.virtual_memory().available / 2**20)) + "Mib" |
nothing calls this directly
no test coverage detected