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
renzhc
diffusers_dcu
Commits
8909ab4b
Unverified
Commit
8909ab4b
authored
May 01, 2024
by
Sayak Paul
Committed by
GitHub
May 01, 2024
Browse files
[Tests] fix: device map tests for models (#7825)
* fix: device module tests * remove patch file * Empty-Commit
parent
c1edb03c
Changes
1
Hide whitespace changes
Inline
Side-by-side
Showing
1 changed file
with
15 additions
and
3 deletions
+15
-3
tests/models/test_modeling_common.py
tests/models/test_modeling_common.py
+15
-3
No files found.
tests/models/test_modeling_common.py
View file @
8909ab4b
...
@@ -691,6 +691,9 @@ class ModelTesterMixin:
...
@@ -691,6 +691,9 @@ class ModelTesterMixin:
def
test_cpu_offload
(
self
):
def
test_cpu_offload
(
self
):
config
,
inputs_dict
=
self
.
prepare_init_args_and_inputs_for_common
()
config
,
inputs_dict
=
self
.
prepare_init_args_and_inputs_for_common
()
model
=
self
.
model_class
(
**
config
).
eval
()
model
=
self
.
model_class
(
**
config
).
eval
()
if
model
.
_no_split_modules
is
None
:
return
model
=
model
.
to
(
torch_device
)
model
=
model
.
to
(
torch_device
)
torch
.
manual_seed
(
0
)
torch
.
manual_seed
(
0
)
...
@@ -718,6 +721,9 @@ class ModelTesterMixin:
...
@@ -718,6 +721,9 @@ class ModelTesterMixin:
def
test_disk_offload_without_safetensors
(
self
):
def
test_disk_offload_without_safetensors
(
self
):
config
,
inputs_dict
=
self
.
prepare_init_args_and_inputs_for_common
()
config
,
inputs_dict
=
self
.
prepare_init_args_and_inputs_for_common
()
model
=
self
.
model_class
(
**
config
).
eval
()
model
=
self
.
model_class
(
**
config
).
eval
()
if
model
.
_no_split_modules
is
None
:
return
model
=
model
.
to
(
torch_device
)
model
=
model
.
to
(
torch_device
)
torch
.
manual_seed
(
0
)
torch
.
manual_seed
(
0
)
...
@@ -728,12 +734,12 @@ class ModelTesterMixin:
...
@@ -728,12 +734,12 @@ class ModelTesterMixin:
model
.
cpu
().
save_pretrained
(
tmp_dir
,
safe_serialization
=
False
)
model
.
cpu
().
save_pretrained
(
tmp_dir
,
safe_serialization
=
False
)
with
self
.
assertRaises
(
ValueError
):
with
self
.
assertRaises
(
ValueError
):
max_size
=
int
(
self
.
model_split_percents
[
1
]
*
model_size
)
max_size
=
int
(
self
.
model_split_percents
[
0
]
*
model_size
)
max_memory
=
{
0
:
max_size
,
"cpu"
:
max_size
}
max_memory
=
{
0
:
max_size
,
"cpu"
:
max_size
}
# This errors out because it's missing an offload folder
# This errors out because it's missing an offload folder
new_model
=
self
.
model_class
.
from_pretrained
(
tmp_dir
,
device_map
=
"auto"
,
max_memory
=
max_memory
)
new_model
=
self
.
model_class
.
from_pretrained
(
tmp_dir
,
device_map
=
"auto"
,
max_memory
=
max_memory
)
max_size
=
int
(
self
.
model_split_percents
[
1
]
*
model_size
)
max_size
=
int
(
self
.
model_split_percents
[
0
]
*
model_size
)
max_memory
=
{
0
:
max_size
,
"cpu"
:
max_size
}
max_memory
=
{
0
:
max_size
,
"cpu"
:
max_size
}
new_model
=
self
.
model_class
.
from_pretrained
(
new_model
=
self
.
model_class
.
from_pretrained
(
tmp_dir
,
device_map
=
"auto"
,
max_memory
=
max_memory
,
offload_folder
=
tmp_dir
tmp_dir
,
device_map
=
"auto"
,
max_memory
=
max_memory
,
offload_folder
=
tmp_dir
...
@@ -749,6 +755,9 @@ class ModelTesterMixin:
...
@@ -749,6 +755,9 @@ class ModelTesterMixin:
def
test_disk_offload_with_safetensors
(
self
):
def
test_disk_offload_with_safetensors
(
self
):
config
,
inputs_dict
=
self
.
prepare_init_args_and_inputs_for_common
()
config
,
inputs_dict
=
self
.
prepare_init_args_and_inputs_for_common
()
model
=
self
.
model_class
(
**
config
).
eval
()
model
=
self
.
model_class
(
**
config
).
eval
()
if
model
.
_no_split_modules
is
None
:
return
model
=
model
.
to
(
torch_device
)
model
=
model
.
to
(
torch_device
)
torch
.
manual_seed
(
0
)
torch
.
manual_seed
(
0
)
...
@@ -758,7 +767,7 @@ class ModelTesterMixin:
...
@@ -758,7 +767,7 @@ class ModelTesterMixin:
with
tempfile
.
TemporaryDirectory
()
as
tmp_dir
:
with
tempfile
.
TemporaryDirectory
()
as
tmp_dir
:
model
.
cpu
().
save_pretrained
(
tmp_dir
)
model
.
cpu
().
save_pretrained
(
tmp_dir
)
max_size
=
int
(
self
.
model_split_percents
[
1
]
*
model_size
)
max_size
=
int
(
self
.
model_split_percents
[
0
]
*
model_size
)
max_memory
=
{
0
:
max_size
,
"cpu"
:
max_size
}
max_memory
=
{
0
:
max_size
,
"cpu"
:
max_size
}
new_model
=
self
.
model_class
.
from_pretrained
(
new_model
=
self
.
model_class
.
from_pretrained
(
tmp_dir
,
device_map
=
"auto"
,
offload_folder
=
tmp_dir
,
max_memory
=
max_memory
tmp_dir
,
device_map
=
"auto"
,
offload_folder
=
tmp_dir
,
max_memory
=
max_memory
...
@@ -774,6 +783,9 @@ class ModelTesterMixin:
...
@@ -774,6 +783,9 @@ class ModelTesterMixin:
def
test_model_parallelism
(
self
):
def
test_model_parallelism
(
self
):
config
,
inputs_dict
=
self
.
prepare_init_args_and_inputs_for_common
()
config
,
inputs_dict
=
self
.
prepare_init_args_and_inputs_for_common
()
model
=
self
.
model_class
(
**
config
).
eval
()
model
=
self
.
model_class
(
**
config
).
eval
()
if
model
.
_no_split_modules
is
None
:
return
model
=
model
.
to
(
torch_device
)
model
=
model
.
to
(
torch_device
)
torch
.
manual_seed
(
0
)
torch
.
manual_seed
(
0
)
...
...
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