From ac4384ef75be8f1ac1eb3ebcfecccf898ccd0c7b Mon Sep 17 00:00:00 2001 From: somebody Date: Wed, 31 May 2023 10:55:46 -0500 Subject: [PATCH] Auto _no_split_modules --- modeling/inference_models/hf_torch.py | 16 +++------------- 1 file changed, 3 insertions(+), 13 deletions(-) diff --git a/modeling/inference_models/hf_torch.py b/modeling/inference_models/hf_torch.py index b00132be..24c0dbb9 100644 --- a/modeling/inference_models/hf_torch.py +++ b/modeling/inference_models/hf_torch.py @@ -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")