[comm][RFC] Implement a custom comm framework#3883
Conversation
meta-cla
Bot
added
the
CLA Signed
Only custom allreduce is supported at this moment. 1. Use torch.library to install the custom comm. 2. Use fixed size-based algorithm to decide which collective to use. 3. Cache one symmetric memory buffer per PG. While this implementation may be torch.compile-able, the current use case is RL only. For torch.compile is enabled, Inductor actually has a path to do a similar replacement. ghstack-source-id: 094f9db Pull-Request: #3883
Only custom allreduce is supported at this moment. 1. Use torch.library to install the custom comm. 2. Use fixed size-based algorithm to decide which collective to use. 3. Cache one symmetric memory buffer per PG. While this implementation may be torch.compile-able, the current use case is RL only. For torch.compile is enabled, Inductor actually has a path to do a similar replacement. ghstack-source-id: dfbb9ee Pull-Request: #3883
tianyu-l
left a comment
tianyu-l
left a comment
There was a problem hiding this comment.
shall we run some benchmarking against vllm AR following
https://github.com/pytorch/torchtitan/blob/main/torchtitan/experiments/rl/.claude/skills/inference_perf_hillclimb/SKILL.md#example-ladder--results
If it's on par we should remove the vllm patch. cc @wwwjn
| view = symm_buffer[: input.numel()].view_as(input) | ||
| out = torch.empty_like(input) | ||
|
|
||
| match algo: |
There was a problem hiding this comment.
What's the consideration of
- putting this dispatcher in torchtitan vs.
- directly modifying pytorch default?
For pytorch side change, is numerics change the main worry?
cc @ngimel
There was a problem hiding this comment.
maybe a weird comment, but would this be a good candidate for an override? I know that only applies to Configurable, but with spmd_types we were planning to make those more configurable anyways, with fwd/bwd dtype control, etc.,
It also already makes sense, to make the src/dst configs collective-type-first (e.g. say "allreduce" instead of placement-based src/dst), and make placements just type assertions, as we won't actually support arbitrary src -> dst redistributes as DTensor does. This could pave the way for easier configurability/overriding
Looking at the overrides doc:
These are valuable for performance but are (a) hardware-specific, (b) harder to maintain, and (c) reasonably held to a lower code-quality bar than core. We want to offer vendors a clean extension point rather than either lowering the core bar or saying "no".
this sounds applicable?
There was a problem hiding this comment.
@pianpwk
Makes sense. My only worry was that creating a config just for patching sounds even hackier than override itself.
But I think it makes sense because people do want to experiment with collectives.
Happy to discuss / review if you could provide a more concrete proposal.
wwwjn
left a comment
wwwjn
left a comment
There was a problem hiding this comment.
It fallback to NCCL comms, so I would assume RL Policytrainer / pretrain will not regression.
So for benchmarking we can focus on generator now (vllm AR patch only patched generator), you can start from
to benchmarking inference only?
|
@fegin do we have benchmarks on the performance here for the different algorithms and data sizes? |
5 participants
Stack from ghstack (oldest at bottom):
Only custom allreduce is supported at this moment.
While this implementation may be torch.compile-able, the current use case is RL only. For torch.compile is enabled, Inductor actually has a path to do a similar replacement.
TODO
Fixed size algorithm is not accurate. Run-time profiling is a better approach but this requires the local PG during initialization, which may not be always created (e.g., if we don't do TP at all).