[Fix][FSDP] Don't remove post backward hooks for multiple backward fix #923

ngoyal2707 · 2022-02-01T19:51:32Z

fixes #918

I am quite confident, that we dont need to remove backward hooks even after finalizing. They will be automatically removed if the leaf variables go out of context and cuda autograd graph cleans up.

Mots tests were succeeding, apart from one related to cpu offload locally, will debug that

ngoyal2707 · 2022-02-01T19:52:51Z

cc: @myleott

min-xu-ai · 2022-02-01T21:09:32Z

I remember if we don't remove it bad things will happen. :-) Exactly which bad things is the good question. I suspect that some backward hooks will fire excessively or some memory will leak. But if we can be confident that this is indeed save to do, it would be really nice! Less code is always better!

ngoyal2707 · 2022-02-01T21:12:56Z

I remember if we don't remove it bad things will happen. :-) Exactly which bad things is the good question. I suspect that some backward hooks will fire excessively or some memory will leak. But if we can be confident that this is indeed save to do, it would be really nice! Less code is always better!

hmm is there any way to test "some backward hooks will fire excessively or some memory will leak"? or any doc or any past issues about this part? or who's the best POC about this? Myle?

min-xu-ai · 2022-02-01T22:45:34Z

hmm is there any way to test "some backward hooks will fire excessively or some memory will leak"? or any doc or any past issues about this part? or who's the best POC about this? Myle?

I am the POC. :-) I remember touch this part of the code last time. I don't remember anyone else touched it after I did.

I wish I remember the details of potential issues with removing it. I think the best way for now is to use the unit tests. After unit tests are good if I can't thinking of any potential issues that are not covered by the unit test, we can certainly go ahead remove these line.

ngoyal2707 · 2022-02-02T20:36:32Z

Updated PR, local tests are still running but I am hoping now they will all pass 🤞

ngoyal2707 · 2022-02-02T21:43:37Z

@min-xu-ai looks like all but MeVo tests are passing, which should be unrelated to this change

min-xu-ai · 2022-02-02T23:16:16Z

@min-xu-ai looks like all but MeVo tests are passing, which should be unrelated to this change

hmm, are you sure? I clicked on this first one, it seems the serialization test failed and then it timed out:

https://app.circleci.com/pipelines/github/facebookresearch/fairscale/4061/workflows/c33d045c-fbe1-4364-8bc7-7a44335cb7d3/jobs/44795

ngoyal2707 · 2022-02-02T23:21:37Z

@min-xu-ai looks like all but MeVo tests are passing, which should be unrelated to this change

hmm, are you sure? I clicked on this first one, it seems the serialization test failed and then it timed out:

https://app.circleci.com/pipelines/github/facebookresearch/fairscale/4061/workflows/c33d045c-fbe1-4364-8bc7-7a44335cb7d3/jobs/44795

that has to be flaky tests 🤷‍♂️ cause all are passing locally and the change has no relation at all to serialization

ngoyal2707 · 2022-02-02T23:21:59Z

(fairseq-20211123-py39-meg-v2.6) namangoyal@learnfair7614:~/src/fairscale$ pytest tests/nn/data_parallel/test_fsdp.py::TestSerialization
=========================================================================================================================================================================== test session starts ============================================================================================================================================================================
platform linux -- Python 3.9.7, pytest-5.4.1, py-1.11.0, pluggy-0.13.1 -- /private/home/namangoyal/.conda/envs/fairseq-20211123-py39-meg-v2.6/bin/python
cachedir: .pytest_cache
rootdir: /private/home/namangoyal/src/fairscale, inifile: setup.cfg
plugins: hydra-core-1.0.7, timeout-1.4.2, cov-2.10.0
collected 8 items

tests/nn/data_parallel/test_fsdp.py::TestSerialization::test_multiprocessing__False_False_ PASSED                                                                                                                                                                                                                                                                    [ 12%]
tests/nn/data_parallel/test_fsdp.py::TestSerialization::test_multiprocessing__False_True_ PASSED                                                                                                                                                                                                                                                                     [ 25%]
tests/nn/data_parallel/test_fsdp.py::TestSerialization::test_multiprocessing__True_False_ PASSED                                                                                                                                                                                                                                                                     [ 37%]
tests/nn/data_parallel/test_fsdp.py::TestSerialization::test_multiprocessing__True_True_ PASSED                                                                                                                                                                                                                                                                      [ 50%]
tests/nn/data_parallel/test_fsdp.py::TestSerialization::test_pickle__False_False_ PASSED                                                                                                                                                                                                                                                                             [ 62%]
tests/nn/data_parallel/test_fsdp.py::TestSerialization::test_pickle__False_True_ PASSED                                                                                                                                                                                                                                                                              [ 75%]
tests/nn/data_parallel/test_fsdp.py::TestSerialization::test_pickle__True_False_ PASSED                                                                                                                                                                                                                                                                              [ 87%]
tests/nn/data_parallel/test_fsdp.py::TestSerialization::test_pickle__True_True_ PASSED                                                                                                                                                                                                                                                                               [100%]

======================================================================================================================================================================= 8 passed in 77.47s (0:01:17) =======================================================================================================================================================================

min-xu-ai · 2022-02-02T23:30:09Z

We do have flaky tests. We try to track them here: #908

However, I just briefly looked at this error:

I have actually never seen it before. But I could be wrong. Let me trigger some re-runs on CI.

min-xu-ai · 2022-02-02T23:32:23Z

@zhaojuanmao, can you take a look at this diff as well? Do you think we could remove the code that removes the hooks?

anj-s · 2022-02-03T01:01:20Z

fairscale/nn/data_parallel/fully_sharded_data_parallel.py

@@ -1543,8 +1543,6 @@ def _register_post_backward_hooks(self) -> None:
            return  # don't register grad hooks if grad isn't enabled
        for p in self.params:
            if p.requires_grad:
-                if hasattr(p, "_shard_bwd_hook"):


Do the comments above need to updated because it seems like we don't need to remove the hook at the end of the BW?

ohh yeah definitely, will do that

if the two lines are removed, will the hooks fire multiple times for multiple forwards cases (e.g. multiple activation checkpointing)?

even for single activation checkpointing case like Checkpoint(FSDP(module)), if there is forward recomputation in the backward pass, the hooks will be registered twice and fired twice unexpectedly?

anj-s · 2022-02-03T01:02:43Z

@ngoyal2707 Can you share the use case that this supports? I don't think I understood the context from the PR title.

ngoyal2707 · 2022-02-03T01:19:52Z

@ngoyal2707 Can you share the use case that this supports? I don't think I understood the context from the PR title.

@anj-s its a fix for this: #918

anj-s · 2022-02-03T01:22:35Z

@ngoyal2707 Can you share the use case that this supports? I don't think I understood the context from the PR title.

@anj-s its a fix for this: #918

Can you add a unit test for this using the smaller repro you have? Great that you figured out the solution!

ngoyal2707 · 2022-02-03T15:40:57Z

@ngoyal2707 Can you share the use case that this supports? I don't think I understood the context from the PR title.

@anj-s its a fix for this: #918

Can you add a unit test for this using the smaller repro you have? Great that you figured out the solution!

yes, definitely, will add unit test as well

min-xu-ai · 2022-02-03T22:29:32Z

This is very strange but I have triggered the re-run twice and both times + the first time failed with exact same set of errors. This could be a CI infra issue. But it doesn't seem to be due to flakiness at this point. :-(

ngoyal2707 · 2022-02-04T02:39:31Z

ime failed with exact same set of errors. This could be a CI infra issue. But it doesn't seem to be due to flaki

hmm thats weird, doesn't reproduces on locally, so unfortunately no idea. how to fix it.

tmarkstrum · 2022-02-04T17:17:02Z

Is there a way to trigger the serialization unit test in the main branch? This can help to identify if the failure is due to the CI infra or the code.

min-xu-ai · 2022-02-04T17:54:57Z

Is there a way to trigger the serialization unit test in the main branch? This can help to identify if the failure is due to the CI infra or the code.

good point. let me trigger a rerun on main

min-xu-ai · 2022-02-04T21:06:37Z

OK. It seems that there is some infra issue with CI. The main branch is failing with timeouts as well

anj-s · 2022-04-04T15:16:38Z

@ngoyal2707 Friendly ping on this PR! Do we still want to check in this fix?

min-xu-ai · 2022-09-24T15:29:58Z

Thanks for this PR! I have taken it over in #1079, with the new test from you as well. Will close this and merge that one in once all tests are done.

facebook-github-bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Feb 1, 2022

ngoyal2707 requested review from anj-s and min-xu-ai February 1, 2022 19:54

anj-s changed the title ~~dont remove post backward hooks for multiple backward fix~~ [Fix][FSDP] Don't remove post backward hooks for multiple backward fix Feb 2, 2022

tmarkstrum self-requested a review February 2, 2022 20:11

dont remove post backward hooks for multiple backward fix

40b9436

ngoyal2707 force-pushed the ngoyal_fix_for_multiple_backwards branch from cc28e10 to 40b9436 Compare February 2, 2022 20:31

anj-s reviewed Feb 3, 2022

View reviewed changes

min-xu-ai mentioned this pull request Sep 24, 2022

[Fix][FSDP] Don't remove post backward hooks for multiple backward fix #1079

Merged

10 tasks

min-xu-ai closed this Sep 24, 2022

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[Fix][FSDP] Don't remove post backward hooks for multiple backward fix #923

[Fix][FSDP] Don't remove post backward hooks for multiple backward fix #923

ngoyal2707 commented Feb 1, 2022

ngoyal2707 commented Feb 1, 2022

min-xu-ai commented Feb 1, 2022

ngoyal2707 commented Feb 1, 2022

min-xu-ai commented Feb 1, 2022

ngoyal2707 commented Feb 2, 2022

ngoyal2707 commented Feb 2, 2022

min-xu-ai commented Feb 2, 2022

ngoyal2707 commented Feb 2, 2022

ngoyal2707 commented Feb 2, 2022

min-xu-ai commented Feb 2, 2022

min-xu-ai commented Feb 2, 2022

anj-s Feb 3, 2022

ngoyal2707 Feb 3, 2022

zhaojuanmao Feb 3, 2022

zhaojuanmao Feb 3, 2022

anj-s commented Feb 3, 2022

ngoyal2707 commented Feb 3, 2022

anj-s commented Feb 3, 2022 •

edited

ngoyal2707 commented Feb 3, 2022

min-xu-ai commented Feb 3, 2022

ngoyal2707 commented Feb 4, 2022

tmarkstrum commented Feb 4, 2022

min-xu-ai commented Feb 4, 2022

min-xu-ai commented Feb 4, 2022

anj-s commented Apr 4, 2022

min-xu-ai commented Sep 24, 2022

[Fix][FSDP] Don't remove post backward hooks for multiple backward fix #923

[Fix][FSDP] Don't remove post backward hooks for multiple backward fix #923

Conversation

ngoyal2707 commented Feb 1, 2022

ngoyal2707 commented Feb 1, 2022

min-xu-ai commented Feb 1, 2022

ngoyal2707 commented Feb 1, 2022

min-xu-ai commented Feb 1, 2022

ngoyal2707 commented Feb 2, 2022

ngoyal2707 commented Feb 2, 2022

min-xu-ai commented Feb 2, 2022

ngoyal2707 commented Feb 2, 2022

ngoyal2707 commented Feb 2, 2022

min-xu-ai commented Feb 2, 2022

min-xu-ai commented Feb 2, 2022

anj-s Feb 3, 2022

Choose a reason for hiding this comment

ngoyal2707 Feb 3, 2022

Choose a reason for hiding this comment

zhaojuanmao Feb 3, 2022

Choose a reason for hiding this comment

zhaojuanmao Feb 3, 2022

Choose a reason for hiding this comment

anj-s commented Feb 3, 2022

ngoyal2707 commented Feb 3, 2022

anj-s commented Feb 3, 2022 • edited

ngoyal2707 commented Feb 3, 2022

min-xu-ai commented Feb 3, 2022

ngoyal2707 commented Feb 4, 2022

tmarkstrum commented Feb 4, 2022

min-xu-ai commented Feb 4, 2022

min-xu-ai commented Feb 4, 2022

anj-s commented Apr 4, 2022

min-xu-ai commented Sep 24, 2022

anj-s commented Feb 3, 2022 •

edited