[rl] Add vllm All-reduce patch in genearator#3759
Conversation
meta-cla
Bot
added
the
CLA Signed
wwwjn
changed the title
|
We implemented this using PyTorch symmetric memory before and it was actually used internally afaik. The implementation even fuse bias addition, which will give you roughly 4 us improvement on small data, iirc. meta-pytorch/kraken#8 |
I briefly read the PR, seems the karken implementation brings roughly the same speedup with vllm's AR at smaller message. I also see vllm has a flag to enable their symmem AR . In my test above ( +21%) , it doesn't use symmetric memory and maybe dispatched to 1 shot AR. I tested enabling symmetric mem, but seems no more gains. n00b question, does symmetric memory needs extra hardware support? If not, can we use karken override in titan? |
| w = self.weight | ||
| w_local = w.to_local() if isinstance(w, DTensor) else w | ||
| y = F.linear(x_local, w_local) | ||
| y = tensor_model_parallel_all_reduce(y) |
There was a problem hiding this comment.
Can we only do ops level dispatch of torchtitan.distributed.All_reduce.
Only generator changes, should be a shorter patch
There was a problem hiding this comment.
I thought about this, one way to do this is to dispatch _AllReduce in spmd_types to TorchTitan's implementation. This will only work if everything is represented in spmd_types. The cost should be cheap.
83fcd67 to
b677c1d
Compare
wwwjn
changed the title
|
We can just copy the kraken implementation to TorchTitan, that would be an easier way to maintain. We wrote the code, and we know how to maintain. Or we can just use vllm's implementation. |
|
|
||
| c10d.all_reduce = all_reduce | ||
|
|
||
| # Force vLLM's TP custom AR onto its registered=False path (see point 2). |
There was a problem hiding this comment.
This part is to avoid vLLM AR to use it's own cuda IPC buffer, which is not compatible with cudagraph.
This would not affect AR performance / speed (verification to be added)
There was a problem hiding this comment.
how is different from torch.ops.symm_mem.one_shot_all_reduce https://docs.pytorch.org/docs/2.12/symmetric_memory.html ?
tianyu-l
left a comment
tianyu-l
left a comment
There was a problem hiding this comment.
please address comments
4 participants
This is clearly not compatible with other parallelism, it's TP only
Test