Add float16 and float32 (#1115)

2023-09-21 00:52:47 -07:00 · 2023-09-21 00:52:47 -07:00 · 1ac4ccf73c
commit 1ac4ccf73c
parent 2ac4d5e2bf
1 changed files with 3 additions and 1 deletions
--- a/vllm/engine/arg_utils.py
+++ b/vllm/engine/arg_utils.py
@ -93,7 +93,9 @@ class EngineArgs:
            '--dtype',
            type=str,
            default=EngineArgs.dtype,
-            choices=['auto', 'half', 'bfloat16', 'float'],
+            choices=[
+                'auto', 'half', 'float16', 'bfloat16', 'float', 'float32'
+            ],
            help='data type for model weights and activations. '
            'The "auto" option will use FP16 precision '
            'for FP32 and FP16 models, and BF16 precision '