Enable autograd graph to propagate after multi-device syncing for loss functions in ddp

[cw-tan]

What does this PR do?

Fixes #2745

Single-line enhancement proposed in #2745, that is, to enable the propagation of the autograd graph after the all_gather operation. This is useful for syncing loss functions in a ddp setting.

Before submitting

• Was this discussed/agreed via a Github issue? (no need for typos and docs improvements)
• Did you read the contributor guideline, Pull Request section?
• Did you make sure to update the docs?
• Did you write any new necessary tests?

PR review

Anyone in the community is free to review the PR once the tests have passed.
If we didn't discuss your PR in Github issues there's a high chance it will not be merged.

Did you have fun?

Make sure you had fun coding 🙃

---

📚 Documentation preview 📚: https://torchmetrics--2754.org.readthedocs.build/en/2754/

[Borda]

That sounds good to me, but can we add a test for this enhancement?

[cw-tan]

That sounds good to me, but can we add a test for this enhancement?

Thanks for the prompt response @Borda.

I'm thinking that _test_ddp_gather_uneven_tensors (here) and _test_ddp_gather_uneven_tensors_multidim (here) in tests/unittests/bases/test_ddp.py already cover the correctness of gather_all_tensors. I'm not sure what other ddp tests there are, but those tests should help tell us if the change I made isn't breaking existing functionality. Let me know if you had something else in mind for this.

I can make an additional unittest in tests/unittests/bases/test_ddp.py to give a tensor that requires_grad to gather_all_tensors, compute some scalar from them (proxy for a loss), and compute grads two ways (one going through the all_gather, one that doesn't) and compare. So this tests that the change achieves the desired effect. How does that sound?

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 69%. Comparing base (abdd2c4) to head (5f29c4d).
Report is 1 commits behind head on master.

Additional details and impacted files

@@           Coverage Diff           @@
##           master   #2754    +/-   ##
=======================================
- Coverage      69%     69%    -0%     
=======================================
  Files         344     330    -14     
  Lines       18824   18653   -171     
=======================================
- Hits        12971   12801   -170     
+ Misses       5853    5852     -1     

[Borda]

I can make an additional unittest in tests/unittests/bases/test_ddp.py to give a tensor that requires_grad to gather_all_tensors, compute some scalar from them (proxy for a loss), and compute grads two ways (one going through the all_gather, one that doesn't) and compare. So this tests that the change achieves the desired effect. How does that sound?

yeah, that sounds good to me :)

[cw-tan]

Update: to accommodate both cases where tensors from different ranks have the same/different shape, the line to put the original tensor (holding the AD graph) back into the gathered list was added in two places in the code.

Because of the two cases, I wrote two unittests to account for each. Interestingly, both pass 2.X stable, but for 1.X LTS, the "same shape" test passes but "different shape" test fails, and for 1.10 oldest, the "different shape" test passes but "same shape" test fails😅. I'll double check for bugs, but the actual code change is just two lines (and all other tests pass, so existing functionality still works), and the unittests are pretty short. The dependency of the unittests passing on different torch versions seems to indicate that it might be a torch versioning issue, maybe to do with ddp behavior? Any thoughts, @Borda ?

[Borda]

I wrote two unittests to account for each. Interestingly, both pass 2.X stable, but for 1.X LTS, the "same shape" test passes but "different shape" test fails, and for 1.10 oldest, the "different shape" test passes but "same shape" test fails😅.

that is strange and worse some more investigation...
cc: @SkafteNicki

[SkafteNicki]

I looked briefly why the tests do not pass on older versions of Pytorch but could not find a reason.

I think we should just only support this for Pytorch > 2.0 and then add this to the documentation.

[SkafteNicki]

seeems the two test functions are now included twice in the test_ddp.py file, please check

[SkafteNicki]

Alright, I finally sat down to understand what was going on here. The non-deterministic behavior was really strange to me, so I tried a lot of debugging and realized that order of output of the all_gather operation was not consistent. This meant that sometimes process 0 was the first element in the gathered list, sometimes the second e.g.

Process 0 gathered tensors: [tensor([0.]), tensor([1.])]  # rank and gathering order matches
Process 1 gathered tensors: [tensor([0.]), tensor([1.])]
Process 0 gathered tensors: [tensor([1.]), tensor([0.])]  # rank and gathering order do not match
Process 1 gathered tensors: [tensor([1.]), tensor([0.])]

I found out that the reason for this is that when the setup function (that initializes the ddp group) is called with another pool.map/pool.starmap than where the test is, this can lead to the processes not matches the expected order. This is exactly what we do in the test code currently

torchmetrics/tests/unittests/conftest.py

Lines 66 to 68 in abdd2c4

	pool = Pool(processes=NUM_PROCESSES)
	pool.starmap(setup_ddp, [(rank, NUM_PROCESSES) for rank in range(NUM_PROCESSES)])
	pytest.pool = pool

and then pytest.pool.starmap is called when we want to run a test in ddp mode. The solution was to call the setup function during the test function and everything is in the expected order. See this commit for details: 48e699b.
This has not been a problem before because we normally reduce all the states in some way e.g. sum them and then the order does not matter at all. Hopefully, this also means that it works regardless of Pytorch version.

@cw-tan sorry for the headache this must have been to debug. I have significantly simplified the tests you had for me to understand what was going on. Hope this still is fine with you.

[cw-tan]

@SkafteNicki Fantastic, thank you so much! I'm just excited to see this feature released so I can remove the monkeypatch in my own code to achieve the same effects. I think the docs are the remaining change -- still some details about only being PyTorch 2 compatible, but the tests have passed for other versions.