Skip to content
GitLab
Menu
Projects
Groups
Snippets
Loading...
Help
Help
Support
Community forum
Keyboard shortcuts
?
Submit feedback
Contribute to GitLab
Sign in / Register
Toggle navigation
Menu
Open sidebar
norm
vllm
Commits
65ea2ddf
Unverified
Commit
65ea2ddf
authored
Nov 16, 2023
by
Aaron Pham
Committed by
GitHub
Nov 16, 2023
Browse files
feat(config): support parsing torch.dtype (#1641)
Signed-off-by:
Aaron
<
29749331+aarnphm@users.noreply.github.com
>
parent
b514d3c4
Changes
1
Hide whitespace changes
Inline
Side-by-side
Showing
1 changed file
with
18 additions
and
12 deletions
+18
-12
vllm/config.py
vllm/config.py
+18
-12
No files found.
vllm/config.py
View file @
65ea2ddf
from
typing
import
Optional
from
typing
import
Optional
,
Union
import
torch
from
transformers
import
PretrainedConfig
...
...
@@ -58,7 +58,7 @@ class ModelConfig:
trust_remote_code
:
bool
,
download_dir
:
Optional
[
str
],
load_format
:
str
,
dtype
:
str
,
dtype
:
Union
[
str
,
torch
.
dtype
]
,
seed
:
int
,
revision
:
Optional
[
str
]
=
None
,
tokenizer_revision
:
Optional
[
str
]
=
None
,
...
...
@@ -331,7 +331,7 @@ _STR_DTYPE_TO_TORCH_DTYPE = {
def
_get_and_verify_dtype
(
config
:
PretrainedConfig
,
dtype
:
str
,
dtype
:
Union
[
str
,
torch
.
dtype
]
,
)
->
torch
.
dtype
:
# NOTE: getattr(config, "torch_dtype", torch.float32) is not correct
# because config.torch_dtype can be None.
...
...
@@ -339,17 +339,23 @@ def _get_and_verify_dtype(
if
config_dtype
is
None
:
config_dtype
=
torch
.
float32
dtype
=
dtype
.
lower
()
if
dtype
==
"auto"
:
if
config_dtype
==
torch
.
float32
:
# Following the common practice, we use float16 for float32 models.
torch_dtype
=
torch
.
float16
if
isinstance
(
dtype
,
str
):
dtype
=
dtype
.
lower
()
if
dtype
==
"auto"
:
if
config_dtype
==
torch
.
float32
:
# Following the common practice, we use float16 for float32
# models.
torch_dtype
=
torch
.
float16
else
:
torch_dtype
=
config_dtype
else
:
torch_dtype
=
config_dtype
if
dtype
not
in
_STR_DTYPE_TO_TORCH_DTYPE
:
raise
ValueError
(
f
"Unknown dtype:
{
dtype
}
"
)
torch_dtype
=
_STR_DTYPE_TO_TORCH_DTYPE
[
dtype
]
elif
isinstance
(
dtype
,
torch
.
dtype
):
torch_dtype
=
dtype
else
:
if
dtype
not
in
_STR_DTYPE_TO_TORCH_DTYPE
:
raise
ValueError
(
f
"Unknown dtype:
{
dtype
}
"
)
torch_dtype
=
_STR_DTYPE_TO_TORCH_DTYPE
[
dtype
]
raise
ValueError
(
f
"Unknown dtype:
{
dtype
}
"
)
# Verify the dtype.
if
torch_dtype
!=
config_dtype
:
...
...
Write
Preview
Markdown
is supported
0%
Try again
or
attach a new file
.
Attach a file
Cancel
You are about to add
0
people
to the discussion. Proceed with caution.
Finish editing this message first!
Cancel
Please
register
or
sign in
to comment