Skip to content
GitLab
Menu
Projects
Groups
Snippets
Loading...
Help
Help
Support
Community forum
Keyboard shortcuts
?
Submit feedback
Contribute to GitLab
Sign in
Toggle navigation
Menu
Open sidebar
OpenDAS
vllm_cscc
Commits
bb76538b
Unverified
Commit
bb76538b
authored
Oct 17, 2024
by
Shashwat Srijan
Committed by
GitHub
Oct 17, 2024
Browse files
[Hardwware][Neuron] Simplify model load for transformers-neuronx library (#9380)
parent
d615b5c9
Changes
1
Show whitespace changes
Inline
Side-by-side
Showing
1 changed file
with
1 addition
and
30 deletions
+1
-30
vllm/model_executor/model_loader/neuron.py
vllm/model_executor/model_loader/neuron.py
+1
-30
No files found.
vllm/model_executor/model_loader/neuron.py
View file @
bb76538b
...
...
@@ -6,7 +6,6 @@ from typing import Dict, List, Optional, Tuple
import
torch
import
torch.nn
as
nn
import
transformers
from
transformers
import
PretrainedConfig
from
vllm.config
import
ModelConfig
,
ParallelConfig
,
SchedulerConfig
...
...
@@ -108,39 +107,11 @@ class NeuronCasualLM(nn.Module):
neuronx_module
=
importlib
.
import_module
(
neuronx_module_path
)
neuronx_model_cls
=
getattr
(
neuronx_module
,
neuronx_model_cls_name
)
split_model_dir
=
f
"
{
model_name_or_path
}
-split"
if
_is_pretrained_neuron_checkpoint
(
model_name_or_path
):
split_model_dir
=
model_name_or_path
elif
not
os
.
path
.
exists
(
f
"
{
model_name_or_path
}
-split"
):
hf_model_cls
=
getattr
(
transformers
,
hf_model_cls_name
)
from
transformers_neuronx.module
import
save_pretrained_split
hf_model
=
hf_model_cls
.
from_pretrained
(
model_name_or_path
,
low_cpu_mem_usage
=
True
)
save_pretrained_split
(
hf_model
,
f
"
{
model_name_or_path
}
-split"
)
self
.
model
=
neuronx_model_cls
.
from_pretrained
(
split_model_dir
,
self
.
model
=
neuronx_model_cls
.
from_pretrained
(
model_name_or_path
,
**
kwargs
)
self
.
model
.
to_neuron
()
def
_is_pretrained_neuron_checkpoint
(
model_name_or_path
:
str
)
->
bool
:
# Checking if the neuron checkpoint is saved in the old format.
if
os
.
path
.
isdir
(
os
.
path
.
join
(
model_name_or_path
,
"pytorch_model.bin"
)):
return
True
# Checking if the neuron checkpoint is saved in the new format.
pretrained_split_files
=
[
"config.json"
,
"generation_config.json"
]
pretrained_split_format
=
".safetensors"
for
file
in
pretrained_split_files
:
file_path
=
os
.
path
.
join
(
model_name_or_path
,
file
)
if
not
os
.
path
.
isfile
(
file_path
):
return
False
for
file
in
os
.
listdir
(
model_name_or_path
):
if
file
.
endswith
(
pretrained_split_format
):
return
True
return
False
def
_get_model_architecture
(
config
:
PretrainedConfig
)
->
str
:
architectures
=
getattr
(
config
,
"architectures"
,
[])
for
arch
in
architectures
:
...
...
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