rixty/KoboldAI-Client
1
1
Fork 0
mirror of https://github.com/KoboldAI/KoboldAI-Client.git synced 2025-06-05 21:59:24 +02:00
Code Issues Projects Releases Wiki Activity

Auto _no_split_modules

Browse Source
This commit is contained in:
somebody
2023-05-31 10:55:46 -05:00
parent 146806eba1
commit ac4384ef75
1 changed files with 3 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files

16
modeling/inference_models/hf_torch.py
View File

@@ -280,26 +280,16 @@ class HFTorchInferenceModel(HFInferenceModel):
try:
model = AutoModelForCausalLM.from_config(self.model_config)
# load_checkpoint_in_model(
# model.model,
# location,
# device_map=device_map
# offload_folder="accelerate-disk-cache",
# dtype="float16",
# offload_state_dict=True
# )
# model.tie_weights()
no_split_module_classes = ["GPTJBlock", "OPTDecoderLayer"]
print("[HUGE SKELETON] MAKING DEVICE MAP")
device_map = infer_auto_device_map(
model,
max_memory={0: "10GiB", 1: "7GiB", "cpu": "15GiB"},
no_split_module_classes=no_split_module_classes,
no_split_module_classes=model._no_split_modules,
dtype="float16",
)
print("[HUGE SKELETON] TYING WEIGHTS")
# TODO: ??
# print("[HUGE SKELETON] TYING WEIGHTS")
model.tie_weights()
print("[HUGE SKELETON] LOADING FROM PRETRAINED")