Keep minimal change.

This commit is contained in:
Light 2023-04-12 23:26:06 +08:00
parent 1d8526849b
commit f3591ccfa1
3 changed files with 53 additions and 22 deletions

View File

@ -238,6 +238,7 @@ Optionally, you can use the following command-line flags:
| `--model_type MODEL_TYPE` | GPTQ: Model type of pre-quantized model. Currently LLaMA, OPT, and GPT-J are supported. |
| `--groupsize GROUPSIZE` | GPTQ: Group size. |
| `--pre_layer PRE_LAYER` | GPTQ: The number of layers to allocate to the GPU. Setting this parameter enables CPU offloading for 4-bit models. |
| `--warmup_autotune` | GPTQ: Enable warmup autotune. Only usable for triton. |
#### FlexGen

View File

@ -61,6 +61,16 @@ def _load_quant(model, checkpoint, wbits, groupsize=-1, faster_kernel=False, exc
model.load_state_dict(safe_load(checkpoint), strict=False)
else:
model.load_state_dict(torch.load(checkpoint), strict=False)
try:
from quant import autotune_warmup, make_quant_attn
# triton branch
make_quant_attn(model)
if shared.args.warmup_autotune:
autotune_warmup(model)
except ImportError: # not triton branch
pass
model.seqlen = 2048
print('Done.')
@ -94,7 +104,26 @@ def load_quantized(model_name):
print("Unknown pre-quantized model type specified. Only 'llama', 'opt' and 'gptj' are supported")
exit()
# Now we are going to try to locate the quantized model file.
# Now we are going to try to locate the quantized model file. I think it's cleaner and supports the new name containing groupsize
path_to_model = Path(f'{shared.args.model_dir}/{model_name}')
pt_path = None
priority_name_list = [
Path(f'{shared.args.model_dir}/{model_name}/{shared.args.wbits}bit-{shared.args.groupsize}g.safetensors'),
Path(f'{shared.args.model_dir}/{model_name}/{shared.args.wbits}bit-{shared.args.groupsize}g.pt'),
Path(f'{shared.args.model_dir}/{model_name}/{shared.args.wbits}bit.safetensors'),
Path(f'{shared.args.model_dir}/{model_name}/{shared.args.wbits}bit.pt'),
Path(f'{shared.args.model_dir}/{model_name}-{shared.args.wbits}bit-{shared.args.groupsize}g.safetensors'),
Path(f'{shared.args.model_dir}/{model_name}-{shared.args.wbits}bit-{shared.args.groupsize}g.pt'),
Path(f'{shared.args.model_dir}/{model_name}-{shared.args.wbits}bit.safetensors'),
Path(f'{shared.args.model_dir}/{model_name}-{shared.args.wbits}bit.pt'),
]
for path in priority_name_list:
if path.exists():
pt_path = path
break
# For compatibility, do we really need this?
if not pt_path:
path_to_model = Path(f'{shared.args.model_dir}/{model_name}')
found_pts = list(path_to_model.glob("*.pt"))
found_safetensors = list(path_to_model.glob("*.safetensors"))

View File

@ -115,6 +115,7 @@ parser.add_argument('--wbits', type=int, default=0, help='GPTQ: Load a pre-quant
parser.add_argument('--model_type', type=str, help='GPTQ: Model type of pre-quantized model. Currently LLaMA, OPT, and GPT-J are supported.')
parser.add_argument('--groupsize', type=int, default=-1, help='GPTQ: Group size.')
parser.add_argument('--pre_layer', type=int, default=0, help='GPTQ: The number of layers to allocate to the GPU. Setting this parameter enables CPU offloading for 4-bit models.')
parser.add_argument('--warmup_autotune', action=argparse.BooleanOptionalAction, default=True, help='GPTQ: Enable warmup autotune. Only usable for triton.')
parser.add_argument('--gptq-bits', type=int, default=0, help='DEPRECATED: use --wbits instead.')
parser.add_argument('--gptq-model-type', type=str, help='DEPRECATED: use --model_type instead.')
parser.add_argument('--gptq-pre-layer', type=int, default=0, help='DEPRECATED: use --pre_layer instead.')