Whitelist 6B in breakmodel

Now that we properly support it, allow the menu option to use breakmodel
This commit is contained in:
henk717 2021-11-27 10:09:54 +01:00 committed by GitHub
parent b56ee07ffa
commit 2bc93ba37a
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 2 additions and 2 deletions

View File

@ -52,7 +52,7 @@ modellist = [
["Load an old GPT-2 model (eg CloverEdition)", "GPT2Custom", ""],
["GPT-Neo 1.3B", "EleutherAI/gpt-neo-1.3B", "8GB"],
["GPT-Neo 2.7B", "EleutherAI/gpt-neo-2.7B", "16GB"],
["GPT-J 6B (HF GIT Required)", "EleutherAI/gpt-j-6B", "24GB"],
["GPT-J 6B", "EleutherAI/gpt-j-6B", "24GB"],
["GPT-2", "gpt2", "1GB"],
["GPT-2 Med", "gpt2-medium", "2GB"],
["GPT-2 Large", "gpt2-large", "4GB"],
@ -354,7 +354,7 @@ if(not vars.model in ["InferKit", "Colab", "OAI", "ReadOnly", "TPUMeshTransforme
import torch
print("{0}Looking for GPU support...{1}".format(colors.PURPLE, colors.END), end="")
vars.hascuda = torch.cuda.is_available()
vars.bmsupported = vars.model in ("EleutherAI/gpt-neo-1.3B", "EleutherAI/gpt-neo-2.7B", "NeoCustom")
vars.bmsupported = vars.model in ("EleutherAI/gpt-neo-1.3B", "EleutherAI/gpt-neo-2.7B", "EleutherAI/gpt-j-6B", "NeoCustom")
if(args.breakmodel is not None and args.breakmodel):
print("WARNING: --breakmodel is no longer supported. Breakmodel mode is now automatically enabled when --layers is used (see --help for details).", file=sys.stderr)
if(args.breakmodel_layers is not None):