Adds reduce_scatter into torchft

[allenwang28]

What does this PR do?

Partially addresses #97 by adding reduce_scatter into torchft.

Concretely, this consists of a few pieces:

• Introducing reduce_scatter into the ProcessGroup following the signature [here](https://github.com/pytorch/pytorch/blob/11f69808c64a65c68a4452250ba7719dcff27c78/torch/csrc/distributed/c10d/PyProcessGroup.hpp#L203
• In ProcessGroup* we essentially follow the behavior of other collectives:
  • In ProcessGroupWrapper, it depends on the parent implementation
  • In ProcessGroupDummy, it writes from the first input into output
  • In ProcessGroupBaby, it asserts inputs and moves underlying storage into shared memory
• Add ReduceScatterOptions in _PickleSafeOptions
• Introduces reduce_scatter as an option in _test_pg, however this necessitated a new function (named _should_run_collective) which was needed as e.g. GLOO does not support reduce_scatter. This function essentially takes the collective, backend and device and copies the logic of the published supported collective matrix.

Tests

Presubmits, and:

$ pytest torchft/process_group_test.py 
============================================= test session starts =============================================
platform linux -- Python 3.10.16, pytest-8.3.4, pluggy-1.5.0
rootdir: /home/allencwang/workspace/torchft
configfile: pytest.ini
plugins: typeguard-2.13.3
collected 16 items                                                                                            

torchft/process_group_test.py ................                                                          [100%]

============================================= 16 passed in 31.44s =============================================
[rank0]:[W206 14:54:24.777939032 CudaIPCTypes.cpp:16] Producer process has been terminated before all shared CUDA tensors released. See Note [Sharing CUDA tensors]

Next steps

The logic of _should_run_collective is a bit confusing, as it allows "non defined backends" like ErrorSwallowing* through, to mimic the old behavior before this change. Testing here could become a bit unwieldy as we add more collectives and so a future step could be to refactor the testing.

One nice change could be to parameterize tests by the collective. This will make potentially failing collectives more explicit and will reduce the time it takes to run individual tests. Likely can do this in the next PR.

[d4l3k]

oh wow -- didn't realize we don't support these on Gloo, good to know! cc @c-p-i-o

[fegin]

ye, we miss many APIs on Gloo.

[wconstab]

this approach seems nice and explicit. but is it possible to instead just try: the test, and except: some specific NYI error? (i'm not sure if we raise a consistent type of NYI exception from backends?)

[allenwang28]

this approach seems nice and explicit. but is it possible to instead just try: the test, and except: some specific NYI error? (i'm not sure if we raise a consistent type of NYI exception from backends?)

Yeah this is a good idea, I modified the block as follows:

    for coll_str, args in collectives:
        try:
            coll = getattr(pg, coll_str)
            work = coll(*args)
            works[coll_str] = work
            work.wait()
            fut = work.get_future()
            fut.wait()
            # Check that all tensor arguments have the expected shapes and dtypes
            check_tensors(args)
        except RuntimeError as e:
            if f"does not support {coll_str}" in str(e):
                # Skip collectives that are not supported by the backend.
                continue
            raise e

[d4l3k]

LGTM, thanks for adding this!

[allenwang28]

Updated the test to be simpler, so I removed the utility functions I previously added. This should remove the need for a test refactor. I wanted to parameterize by collective, but #103 shows that tests got much slower after doing this. I will deprecate #103.

I have also added an explicit NotImplementedError for reduce_scatter within ProcessGroupBabyGloo, because otherwise the exception lives within the tx queue (see here). Noticed this as test_baby_gloo_apis would fail here, and that would be a nasty issue for a downstream user. Therefore, adding this explicitly in the API is ultimately cleaner.

[d4l3k]

LGTM