Add `keep_torch_compile` param to `unwrap_model` and `extract_model_from_parallel` for distributed compiled model. #3282

ggoggam · 2024-12-09T06:55:06Z

What does this PR do?

This PR fixes the unexpected behavior of Accelerator.unwrap_model (Issue #3281). Right now, if the model is wrapped in both distributed wrapper (e.g. DistributedDataParallel or DeepSpeedEngine) and compiled module (OptimizedModule) it only unwraps the distributed module. This behavior arises from the following code in L80 of utils/others.py:

accelerate/src/accelerate/utils/other.py

Line 80 in cb8b7c6

is_compiled = is_compiled_module(model)

Instead of checking for compiled model both before and after unwrapping distributed wrapper, the current code only checks for compilation before unwrapping the distributed wrapper. If the model is wrapped in both, is_compiled will be set to False and won't unwrap the model fully, resulting in unexpected behavior (users expect fully unwrapped model before saving, but gets OptimizedModule instead, which may result in an error when loading the state dict due to key mismatch).

Before submitting

This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
Did you read the contributor guideline,
Pull Request section?
Was this discussed/approved via a Github issue or the forum? Please add a link
to it if that's the case.
Did you make sure to update the documentation with your changes? Here are the
documentation guidelines, and
here are tips on formatting docstrings.
Did you write any new necessary tests?

Related Issue

Accelerate.unwrap should fully unwrap the model. #3281

Who can review?

@muellerzr @SunMarc

SunMarc

Hi @ggoggam, this is actually intended behavior as you can see from this PR #1437.

What I can suggest modify either unwrap_model or add a new argument to extract_model_from_parallel. cc @muellerzr

ggoggam · 2024-12-09T14:48:30Z

I see. I think it would make more sense to modify unwrap_model to the exhibit expected behavior. I will wait for other reviews before I implement the fix.

BenjaminBossan · 2024-12-10T14:32:28Z

Thanks for the PR. Regarding the implementation, extract_model_from_parallel is already supposed to handle both compilation and models wrapped by DS/DDP. If it doesn't work, I'd rather see extract_model_from_parallel fixed than calling it repeatedly. Also, could you please add a unit test similar to this one, your original example could be adjusted for that.

ggoggam · 2024-12-11T00:25:24Z

Thanks for the PR. Regarding the implementation, extract_model_from_parallel is already supposed to handle both compilation and models wrapped by DS/DDP. If it doesn't work, I'd rather see extract_model_from_parallel fixed than calling it repeatedly. Also, could you please add a unit test similar to this one, your original example could be adjusted for that.

My first commit actually fixes extract_model_from_parallel and the unit test. While I was fixing this, I noticed that the unit test you mentioned checks for the inner module of the unwrapped model, which is not expected.

accelerate/tests/test_utils.py

Line 252 in cb8b7c6

assert compiled_model._orig_mod == compiled_model_unwrapped._orig_mod

If I understand correctly, it should be

assert compiled_model._orig_mod == compiled_model_unwrapped

if extract_model_from_parallel is supposed to handle both compilation and models by DS/DDP?

BenjaminBossan · 2024-12-11T11:35:22Z

If I understand correctly, it should be
assert compiled_model._orig_mod == compiled_model_unwrapped
if extract_model_from_parallel is supposed to handle both compilation and models by DS/DDP?

Hmm, good question, I'll leave that to the others to answer, as I'm not sure.

SunMarc · 2024-12-11T13:36:03Z

If I understand correctly, it should be
assert compiled_model._orig_mod == compiled_model_unwrapped
if extract_model_from_parallel is supposed to handle both compilation and models by DS/DDP?

That's right. But the current behavior is that we don't unwrap the compiled model with extract_model_from_parallel. Right now, we just make sure that we are able to unwrap the distributed part even if the model is compiled.

cc @muellerzr what do you prefer, modify how extract_model_from_parallel works or modify unwrap function ?

ggoggam · 2024-12-11T14:07:27Z

src/accelerate/utils/other.py

@@ -77,8 +77,7 @@ def extract_model_from_parallel(model, keep_fp32_wrapper: bool = True, recursive
    """
    options = (torch.nn.parallel.DistributedDataParallel, torch.nn.DataParallel)

-    is_compiled = is_compiled_module(model)


To clarify: Is is_compiled being assigned here and then checked after DDP/DS is unwrapped? If so, this seems potentially unnecessary or wrong as it checks for compilation only before unwrapping DDP/DS

Yes I agree that we should check for compilation after

Sorry I think I misunderstood the code. is_compiled is assigned here before unwrapping to keep the torch compilation, though I still think this could be a problem if the model is in the form of Distributed(Compiled(model)). Refer to my newest commit.

muellerzr

Hey! Thanks for trying to tackle this. I think what I'd rather see instead is in extract_model_from_parallel we should add a new param arg called keep_torch_compile, which similar to keep_fp32_wrapper should default to False for a number of versions and then we likely should flip that to True after awhile, & modify the logic in extract_model_from_parallel to reflect this change.

Can you tweak this PR to so? :)

ggoggam · 2024-12-11T16:12:45Z

Hey! Thanks for trying to tackle this. I think what I'd rather see instead is in extract_model_from_parallel we should add a new param arg called keep_torch_compile, which similar to keep_fp32_wrapper should default to False for a number of versions and then we likely should flip that to True after awhile, & modify the logic in extract_model_from_parallel to reflect this change.

Can you tweak this PR to so? :)

Sure thing. I added keep_torch_compile argument to both Accelerator.unwrap_model (defaults to False) and extract_model_from_parallel (defaults to True). Let me know if this makes sense.

I am also curious about what you think in the case of Distributed(Compiled(model)), since the current code only accounts for Compiled(Distributed(model)) case. I think Accelerator.prepare does the latter, but I am not sure if we should still consider the former case.

SunMarc

Thanks for adding this and the tests !
Also, note that we have to also upstream the modification done to extract_model_from_parallel in transformers if needed.

src/accelerate/utils/other.py

HuggingFaceDocBuilderDev · 2024-12-13T14:48:12Z

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

src/accelerate/utils/other.py

SunMarc · 2024-12-13T15:15:35Z

src/accelerate/accelerator.py

@@ -2601,7 +2601,7 @@ def pad_across_processes(self, tensor, dim=0, pad_index=0, pad_first=False):
        """
        return pad_across_processes(tensor, dim=dim, pad_index=pad_index, pad_first=pad_first)

-    def unwrap_model(self, model, keep_fp32_wrapper: bool = True):
+    def unwrap_model(self, model, keep_fp32_wrapper: bool = True, keep_torch_compile: bool = False):


I think you wanted to set keep_torch_compile to True for a couple of version @muellerzr ?

Correct, let's default to True until 2.5.0, to give users ~3 months.

Essentially we should default to None, and then if it gets None warn that the default to this will be changing

I've defaulted keep_torch_compile=True in the latest commit.

…odel_from_parallel`.

ggoggam · 2024-12-24T02:38:05Z

@SunMarc Please let me know if there is any changes needed for this PR!

SunMarc

LGTM !

SunMarc · 2024-12-24T11:40:59Z

cc @muellerzr for final

ggoggam force-pushed the bugfix/check-compiled-model branch from fb3809e to f601b8c Compare December 9, 2024 07:03

SunMarc reviewed Dec 9, 2024

View reviewed changes

ggoggam changed the title ~~Fix extract_model_from_parallel to fully unwrap compiled model.~~ Fix unwrap_model for distributed compiled model. Dec 9, 2024

ggoggam commented Dec 11, 2024

View reviewed changes

muellerzr reviewed Dec 11, 2024

View reviewed changes

ggoggam requested a review from muellerzr December 12, 2024 23:24

SunMarc approved these changes Dec 13, 2024

View reviewed changes

src/accelerate/utils/other.py Show resolved Hide resolved

SunMarc reviewed Dec 13, 2024

View reviewed changes

src/accelerate/utils/other.py Outdated Show resolved Hide resolved

SunMarc reviewed Dec 13, 2024

View reviewed changes

Adding keep_torch_compile argument to unwrap_model and `extract_m…

c2e9e5b

…odel_from_parallel`.

ggoggam force-pushed the bugfix/check-compiled-model branch from 71d8bf1 to c2e9e5b Compare December 18, 2024 09:10

ggoggam requested a review from SunMarc December 19, 2024 07:55

ggoggam changed the title ~~Fix unwrap_model for distributed compiled model.~~ Add keep_torch_compile param to unwrap_model and extract_model_from_parallel for distributed compiled model. Dec 19, 2024

SunMarc approved these changes Dec 24, 2024

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add `keep_torch_compile` param to `unwrap_model` and `extract_model_from_parallel` for distributed compiled model. #3282

Add `keep_torch_compile` param to `unwrap_model` and `extract_model_from_parallel` for distributed compiled model. #3282

ggoggam commented Dec 9, 2024 •

edited

Loading

SunMarc left a comment

ggoggam commented Dec 9, 2024 •

edited

Loading

BenjaminBossan commented Dec 10, 2024

ggoggam commented Dec 11, 2024 •

edited

Loading

BenjaminBossan commented Dec 11, 2024

SunMarc commented Dec 11, 2024

ggoggam Dec 11, 2024 •

edited

Loading

muellerzr Dec 11, 2024

ggoggam Dec 11, 2024 •

edited

Loading

muellerzr left a comment

ggoggam commented Dec 11, 2024

SunMarc left a comment •

edited

Loading

HuggingFaceDocBuilderDev commented Dec 13, 2024

SunMarc Dec 13, 2024

muellerzr Dec 13, 2024

ggoggam Dec 17, 2024

ggoggam commented Dec 24, 2024

SunMarc left a comment

SunMarc commented Dec 24, 2024

Add keep_torch_compile param to unwrap_model and extract_model_from_parallel for distributed compiled model. #3282

Are you sure you want to change the base?

Add keep_torch_compile param to unwrap_model and extract_model_from_parallel for distributed compiled model. #3282

Conversation

ggoggam commented Dec 9, 2024 • edited Loading

What does this PR do?

Before submitting

Related Issue

Who can review?

SunMarc left a comment

Choose a reason for hiding this comment

ggoggam commented Dec 9, 2024 • edited Loading

BenjaminBossan commented Dec 10, 2024

ggoggam commented Dec 11, 2024 • edited Loading

BenjaminBossan commented Dec 11, 2024

SunMarc commented Dec 11, 2024

ggoggam Dec 11, 2024 • edited Loading

Choose a reason for hiding this comment

muellerzr Dec 11, 2024

Choose a reason for hiding this comment

ggoggam Dec 11, 2024 • edited Loading

Choose a reason for hiding this comment

muellerzr left a comment

Choose a reason for hiding this comment

ggoggam commented Dec 11, 2024

SunMarc left a comment • edited Loading

Choose a reason for hiding this comment

HuggingFaceDocBuilderDev commented Dec 13, 2024

SunMarc Dec 13, 2024

Choose a reason for hiding this comment

muellerzr Dec 13, 2024

Choose a reason for hiding this comment

ggoggam Dec 17, 2024

Choose a reason for hiding this comment

ggoggam commented Dec 24, 2024

SunMarc left a comment

Choose a reason for hiding this comment

SunMarc commented Dec 24, 2024

Add `keep_torch_compile` param to `unwrap_model` and `extract_model_from_parallel` for distributed compiled model. #3282

Add `keep_torch_compile` param to `unwrap_model` and `extract_model_from_parallel` for distributed compiled model. #3282

ggoggam commented Dec 9, 2024 •

edited

Loading

ggoggam commented Dec 9, 2024 •

edited

Loading

ggoggam commented Dec 11, 2024 •

edited

Loading

ggoggam Dec 11, 2024 •

edited

Loading

ggoggam Dec 11, 2024 •

edited

Loading

SunMarc left a comment •

edited

Loading