AI/text-generation-webui
1
0
Fork 0
mirror of https://github.com/oobabooga/text-generation-webui.git synced 2024-10-01 01:26:03 -04:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #17 from Silver267/main

Browse Source 
Add --disk-cache-dir and --max-cpu-mem parameters (Properly pushed to github)
This commit is contained in:
oobabooga 2023-01-21 00:26:19 -03:00 committed by GitHub
commit 86a2832f3b
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 12 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
README.md
View File

@ -139,7 +139,9 @@ Optionally, you can use the following command-line flags:
| `--load-in-8bit` | Load the model with 8-bit precision.| | `--load-in-8bit` | Load the model with 8-bit precision.|
| `--auto-devices` | Automatically split the model across the available GPU(s) and CPU.| | `--auto-devices` | Automatically split the model across the available GPU(s) and CPU.|
| `--disk` | If the model is too large for your GPU(s) and CPU combined, send the remaining layers to the disk. | | `--disk` | If the model is too large for your GPU(s) and CPU combined, send the remaining layers to the disk. |
| `--disk-cache-dir DISK_CACHE_DIR` | If you want to specify the drive or folder for offloading to disk, specify the **full path** of your folder here. (Unless you just want to rename your cache folder) |
| `--gpu-memory GPU_MEMORY` | Maximum memory in GiB to allocate to the GPU when loading the model. This is useful if you get out of memory errors while trying to generate text. Must be an integer number. | | `--gpu-memory GPU_MEMORY` | Maximum memory in GiB to allocate to the GPU when loading the model. This is useful if you get out of memory errors while trying to generate text. Must be an integer number. |
| `--max-cpu-mem MAX_CPU_MEMORY` | Maximum memory in GiB to allocate to the CPU when offloading the model to RAM. This parameter defaults to 99GiB. |
| `--no-stream` | Don't stream the text output in real time. This slightly improves the text generation performance.| | `--no-stream` | Don't stream the text output in real time. This slightly improves the text generation performance.|
| `--settings SETTINGS_FILE` | Load the default interface settings from this json file. See `settings-template.json` for an example.| | `--settings SETTINGS_FILE` | Load the default interface settings from this json file. See `settings-template.json` for an example.|
| `--listen` | Make the web UI reachable from your local network.| | `--listen` | Make the web UI reachable from your local network.|

8
server.py
View File

@ -31,6 +31,8 @@ parser.add_argument('--no-stream', action='store_true', help='Don\'t stream the
parser.add_argument('--settings', type=str, help='Load the default interface settings from this json file. See settings-template.json for an example.') parser.add_argument('--settings', type=str, help='Load the default interface settings from this json file. See settings-template.json for an example.')
parser.add_argument('--listen', action='store_true', help='Make the web UI reachable from your local network.') parser.add_argument('--listen', action='store_true', help='Make the web UI reachable from your local network.')
parser.add_argument('--share', action='store_true', help='Create a public URL. This is useful for running the web UI on Google Colab or similar.') parser.add_argument('--share', action='store_true', help='Create a public URL. This is useful for running the web UI on Google Colab or similar.')
parser.add_argument('--max-cpu-mem', type=int, help='Maximum cpu memory in GiB to allocate to the memory for offloading.')
parser.add_argument('--disk-cache-dir', type=str, help='Directory which you want the disk cache to load to.')
args = parser.parse_args() args = parser.parse_args()
loaded_preset = None loaded_preset = None
@ -90,8 +92,14 @@ def load_model(model_name):
else: else:
settings.append("device_map='auto'") settings.append("device_map='auto'")
if args.gpu_memory is not None: if args.gpu_memory is not None:
if args.max_cpu_mem is not None:
settings.append(f"max_memory={{0: '{args.gpu_memory}GiB', 'cpu': '{args.max_cpu_mem}GiB'}}")
else:
settings.append(f"max_memory={{0: '{args.gpu_memory}GiB', 'cpu': '99GiB'}}") settings.append(f"max_memory={{0: '{args.gpu_memory}GiB', 'cpu': '99GiB'}}")
if args.disk: if args.disk:
if args.disk_cache_dir is not None:
settings.append("offload_folder='"+args.disk_cache_dir+"'")
else:
settings.append("offload_folder='cache'") settings.append("offload_folder='cache'")
if args.load_in_8bit: if args.load_in_8bit:
settings.append("load_in_8bit=True") settings.append("load_in_8bit=True")