support async in nccl pg

[tushar00jain]

Summary:

• set the same stream as the one used for work in future continuations so that random streams don't depend on pg stream (this can make these streams dependent on the allreduce stream)
• wait on the work sent to pg's immediately on the fragment streams (used for allreduce) to make them depend on the pg stream and so that they don't depend on any future work that's submitted to those streams
• copy grads before allreduce so that the inner optimization can use those and it doesn't create a dependency between the default stream and the pg stream
• add back support for quantized allreduce in manager
• change return types to be consistent with pg allreduce
• the returned future from quantization collectives hangs (likely because set_result is not called?) so changed it to return the future directly from the pg

Test Plan:

• tested the changes with nccl pg
• synchronize on recovery stream sometimes makes the cpu block on collective (probably because some callback gets scheduled on the recovery stream? we need to remove synchronizing on recovery stream when there is no need to)
• calling work.wait returned by baby nccl pg makes the cpu block on the collective (because 2 contexts can't overlap?)
• pg gloo needs us to call future.wait in the sync phase instead of the prepare phase, so we probably need a different wrapper
• same for baby gloo pg

Without Quantization

With Quantization

---

Stack created with Sapling. Best reviewed with ReviewStack.

• enable merging parameters for diloco #212
• add debugability for baby pg #213
• -> support async in nccl pg #211

[d4l3k]

These changes generally seem reasonable though as always with streams it's a bit hard to follow the logic. It would be nice to add some tests but I'm not sure how to unit test stream logic.

Maybe we could write a stream mock and then check for certain calls/ops but don't think that would play well with the normal operations unfortunately. I think we can land this as is unless you have some ideas

[d4l3k]

if we call work.wait() there's no need for get_future, just run the callback inline

[tushar00jain]

we're getting the fut here so that we can return it

[tushar00jain]

These changes generally seem reasonable though as always with streams it's a bit hard to follow the logic. It would be nice to add some tests but I'm not sure how to unit test stream logic.

Maybe we could write a stream mock and then check for certain calls/ops but don't think that would play well with the normal operations unfortunately. I think we can land this as is unless you have some ideas

@d4l3k yeah that'll be something good to have but likely to be a lot of code so would prefer doing it separately. mocking will also need us to assume a lot of things about streams, which may or may not be correct. this is why looking at profiles has been useful. i have a bunch of feature asks from the profiler :)

[H-Huang]

For my understanding, is what was happening before:

1. computation stream
2. (torch.distributed) collective stream
3. diloco fragment stream (many)

1 was waiting on 2 which waited for 3

Now after the PR:

torchft does not queue any collectives on 2, 1 does not wait for 3

[tushar00jain]

For my understanding, is what was happening before:

1. computation stream
2. (torch.distributed) collective stream
3. diloco fragment stream (many)

1 was waiting on 2 which waited for 3

Now after the PR:

torchft does not queue any collectives on 2, 1 does not wait for 3

@H-Huang you mean because of the copied out gradients?

[H-Huang]

i mean after the stream changes, what was the changed behavior for the streams and what stream is the torchft allreduce running on?

[tushar00jain]

@H-Huang there's also a 4th stream involved that runs the callbacks

1. computation stream
2. (torch.distributed) collective stream
3. diloco fragment stream (many)
4. callback stream
5. recovery stream (sorry 1 more)

i think this is how it works

Before

• 3 waits for 2
• 4 waits for 2
• 3 waits for 4
• we sync 3' (for another fragment)
• we sync 1

3 will block cpu if 4 == 3' or 4 == 1

After

we enforce 4 == 3, so

• 3 waits for 2
• we sync 3'
• we sync 1

but i'm not sure now how the 5 is getting involved since 4 != 5