Functional collectives

mayank31398: <p>Hi, I recently noticed this line in the all-reduce and broadcast collectives<br>
Do we need to do these redundant copies?<br>
This is making TP with DTensors slower than TP with normal tensors.</p>
<pre data-code-wrap="python"><code class="lang-python">at::Tensor all_reduce(
    const at::Tensor&amp; input,
    std::string reduce_op,
    std::string group_name) {
  auto output = input.clone(at::MemoryFormat::Contiguous);
  return all_reduce_(output, std::move(reduce_op), std::move(group_name));
}
</code></pre>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/4/9/49fe1be828480ff25bb5b1e7b115d9981eab7653.png" data-download-href="https://discuss.pytorch.org/uploads/default/49fe1be828480ff25bb5b1e7b115d9981eab7653" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/4/9/49fe1be828480ff25bb5b1e7b115d9981eab7653_2_690x79.png" alt="image" data-base62-sha1="ayzhrEJob1x4XvN99OUY8MND2Lh" width="690" height="79" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/4/9/49fe1be828480ff25bb5b1e7b115d9981eab7653_2_690x79.png, https://discuss.pytorch.org/uploads/default/optimized/3X/4/9/49fe1be828480ff25bb5b1e7b115d9981eab7653_2_1035x118.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/4/9/49fe1be828480ff25bb5b1e7b115d9981eab7653_2_1380x158.png 2x" data-dominant-color="E1D4EA"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">2244×260 23.5 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>

mayank31398: <p>every all reduce call is causing this tensor copy as can be seen in the trace</p>

agu: <p>I think functional all-reduce is always out-of-place, so it does this clone. I think <code>torch.compile</code> would help re-inplace it. cc: <a class="mention" href="/u/wanchaol">@wanchaol</a></p>

wanchaol: <p>Yes <a class="mention" href="/u/agu">@agu</a> is correct, functional collectives by design is out of place so it would always allocating the output tensor first, if turning on torch.compile it would re-inplace it.</p>
<p><a class="mention" href="/u/mayank31398">@mayank31398</a> I wonder how much slowness you observed?</p>
<p>For training we usually use SequenceParallel by default, where for SequenceParallel the allocation would happen anyways and the input/output shape of the allgather/reduce_scatter are different. For the case where there’s no SequenceParallel, iirc we benchmarked the e2e training on llama models and it does not show observable slowness compare to a TP that works on normal tensors.</p>

mayank31398: <p><a class="mention" href="/u/wanchaol">@wanchaol</a> there is a reasonable difference when not using compile.<br>
Is there not a plan of moving to in-place?<br>
The copy is redundant and honestly compile doesn’t work with everything so there are a lot of problems with using out-of-place here</p>

tianyu: <p>continuing discussions in <a href="https://github.com/pytorch/pytorch/issues/134388" class="inline-onebox" rel="noopener nofollow ugc">Remove redundant copy in functional collectives · Issue #134388 · pytorch/pytorch · GitHub</a></p>

syed-ahmed: <p><a class="mention" href="/u/mayank31398">@mayank31398</a> Can we do TP with normal tensors in PyTorch? I thought TP can only be done with DTensors in PyTorch?</p>

mayank31398: <p>yes ofc, why do you think we can’t do TP without DTensors?<br>
following code does TP without sequence parallel (SP):</p>
<pre data-code-wrap="python"><code class="lang-python">class ReduceFromTensorParallelRegion(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x: torch.Tensor) -&gt; torch.Tensor:
        torch.distributed.all_reduce(x)
        return x

    @staticmethod
    def backward(ctx: Any, x_grad: torch.Tensor) -&gt; Any:
        return x_grad
</code></pre>