Skip to content

Enable multi-device for efficientnet#29989

Merged
amyeroberts merged 1 commit into
huggingface:mainfrom
jla524:efficientnet_multidevice
Apr 3, 2024
Merged

Enable multi-device for efficientnet#29989
amyeroberts merged 1 commit into
huggingface:mainfrom
jla524:efficientnet_multidevice

Conversation

@jla524

@jla524 jla524 commented Apr 2, 2024

Copy link
Copy Markdown
Contributor

What does this PR do?

Fixes #29786 (issue)

Tested on two systems, one with 2x RTX 3060 and one with 2x RTX 3090

$ pytest tests/models/efficientnet/test_modeling_efficientnet.py -k "parallel or offload"
=================================================================== test session starts ===================================================================
platform linux -- Python 3.10.13, pytest-7.4.4, pluggy-1.4.0
rootdir: /repos/transformers
configfile: pyproject.toml
plugins: hypothesis-6.98.10, xdist-3.5.0, timeout-2.3.1, anyio-4.3.0
collected 108 items / 101 deselected / 7 selected                                                                                                         

tests/models/efficientnet/test_modeling_efficientnet.py .......                                                                                     [100%]

<warnings redacted>
===================================================== 7 passed, 101 deselected, 12 warnings in 8.85s ======================================================

Who can review?

@amyeroberts

config_class = EfficientNetConfig
base_model_prefix = "efficientnet"
main_input_name = "pixel_values"
_no_split_modules = []

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
_no_split_modules = []
_no_split_modules = ["EfficientNetBlock"]

any reason not to set a correct module to not split?

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@ArthurZucker I don't think it's necessary for there to be a module defined - this is the case for some of our models already in the library e.g. Camembert.

It is strange not defining "EfficientNetBlock" is OK though, as the block uses a residual connection, which requires the two tensors to be on the same device (as then so too the weights).

@ArthurZucker ArthurZucker Apr 2, 2024

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Alright, it's not necessary, but for small GPU it helps. We'll see how it goes

@amyeroberts
amyeroberts merged commit 03732de into huggingface:main Apr 3, 2024
@jla524
jla524 deleted the efficientnet_multidevice branch April 12, 2024 05:26
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Community contribution: enabling device_map="auto" support for more vision and multimodal models

3 participants