Allow for force unwrapping #2595
Conversation
| @@ -193,6 +206,31 @@ def test_extract_model(self): | |||
|
|
|||
| assert model == model_unwrapped | |||
|
|
|||
| @require_tpu | |||
| @require_huggingface_suite | |||
| def test_extract_model_recursive_fsdpv2(self): | |||
There was a problem hiding this comment.
While we don't have runners for this test, @shub-kris verified it works
One more thing that needs to be in the test which I missed is: One has to make sure that you set this two environment variables:
export PJRT_DEVICE=TPU XLA_USE_SPMD=1
I encountered the error when I didn't set XLA_USE_SPMD=1
Please enable SPMD via `torch_xla.runtime.use_spmd()`
There was a problem hiding this comment.
I will do a test-run on this huggingface/transformers#29659 to make sure training and loading of model works properly.
There was a problem hiding this comment.
@muellerzr shouldn't we add that in the test function that it should be ran with
export PJRT_DEVICE=TPU XLA_USE_SPMD=1There was a problem hiding this comment.
I will do a test-run on this huggingface/transformers#29659 to make sure training and loading of model works properly.
Ahh I can't do as the changes aren't there yet in the transformers library.
|
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. |
There was a problem hiding this comment.
Overall LGTM to me thanks. I only think it could be a bit better commented for future clarity.
Regarding the testing, good to have confirmation that this fixes the issue. I guess we could come up with an artificial test case that has module.module and works without TPU. Not sure how useful that would be.
| if recursive: | ||
| # This is needed in cases such as using FSDPv2 on XLA |
There was a problem hiding this comment.
Is there a way to detect that we are using fsdpv2 or xla , so that we don't have to set recursive. cc @shub-kris
There was a problem hiding this comment.
There are, but accelerate doesn’t support it yet and it’s experimental. So for now a generic approach, as it’s a single location in the Trainer
There was a problem hiding this comment.
Thanks for the nice PR @muellerzr. Great work.
I will approve the PR once I do a small training to verify everything works end-to-end.
What does this PR do?
There are cases when we need to unwrap far more than just what is done natively in Accelerate, such as when we have TPUs/SPMD.
This PR expands
extract_model_from_parallelto support situations such as using XLA'sFSDPv2.Alternative to huggingface/transformers#29780, which we can upstream to transformers by just using
recursive=TrueBefore submitting
Pull Request section?
to it if that's the case.
documentation guidelines, and
here are tips on formatting docstrings.
Who can review?
Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.
@SunMarc @shub-kris
We can't test it on TPUs since we don't have our runners, but @shub-kris tested this himself :)