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

Add --panic to raise when loading fails

Browse Source
This commit is contained in:
somebody
2023-07-08 14:04:46 -05:00
parent fd6f66a98d
commit c2ee30af32
3 changed files with 8 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
aiserver.py
View File

@@ -1400,6 +1400,7 @@ def general_startup(override_args=None):
parser.add_argument('-f', action='store', help="option for compatability with colab memory profiles") parser.add_argument('-f', action='store', help="option for compatability with colab memory profiles")
parser.add_argument('-v', '--verbosity', action='count', default=0, help="The default logging level is ERROR or higher. This value increases the amount of logging seen in your screen") parser.add_argument('-v', '--verbosity', action='count', default=0, help="The default logging level is ERROR or higher. This value increases the amount of logging seen in your screen")
parser.add_argument('-q', '--quiesce', action='count', default=0, help="The default logging level is ERROR or higher. This value decreases the amount of logging seen in your screen") parser.add_argument('-q', '--quiesce', action='count', default=0, help="The default logging level is ERROR or higher. This value decreases the amount of logging seen in your screen")
parser.add_argument("--panic", action='store_true', help="Disables falling back when loading fails.")
#args: argparse.Namespace = None #args: argparse.Namespace = None
if "pytest" in sys.modules and override_args is None: if "pytest" in sys.modules and override_args is None:

2
modeling/inference_models/generic_hf_torch/class.py
View File

@@ -90,6 +90,8 @@ class model_backend(HFTorchInferenceModel):
utils.module_names = list(metamodel.state_dict().keys()) utils.module_names = list(metamodel.state_dict().keys())
utils.named_buffers = list(metamodel.named_buffers(recurse=True)) utils.named_buffers = list(metamodel.named_buffers(recurse=True))
except Exception as e: except Exception as e:
if utils.args.panic:
raise e
logger.warning(f"Gave up on lazy loading due to {e}") logger.warning(f"Gave up on lazy loading due to {e}")
self.lazy_load = False self.lazy_load = False

5
modeling/inference_models/hf_torch.py
View File

@@ -363,6 +363,8 @@ class HFTorchInferenceModel(HFInferenceModel):
return GPTNeoForCausalLM.from_pretrained(location, **tf_kwargs) return GPTNeoForCausalLM.from_pretrained(location, **tf_kwargs)
except Exception as e: except Exception as e:
logger.warning(f"{self.model_name} is a no-go; {e} - Falling back to auto.") logger.warning(f"{self.model_name} is a no-go; {e} - Falling back to auto.")
if utils.args.panic:
raise e
# Try to determine model type from either AutoModel or falling back to legacy # Try to determine model type from either AutoModel or falling back to legacy
try: try:
@@ -414,6 +416,9 @@ class HFTorchInferenceModel(HFInferenceModel):
logger.error("Invalid load key! Aborting.") logger.error("Invalid load key! Aborting.")
raise raise
if utils.args.panic:
raise e
logger.warning(f"Fell back to GPT2LMHeadModel due to {e}") logger.warning(f"Fell back to GPT2LMHeadModel due to {e}")
logger.debug(traceback.format_exc()) logger.debug(traceback.format_exc())