_functional_collectives used in torch.compile

yang_xiaorun: <ol>
<li>I’m confused why allreduce is not in traceable_collective_remaps. Because I found that since reduce_scatter_tensor is in traceable_collective_remaps, I can use torch.distributed.distributed_c10d.reduce_scatter_tensor in the script, torch.compile works. torch.distributed.all_reduce cannot be used directly in this way, but is it planned to support torch.distributed.all_reduce in the future?</li>
</ol>
<pre><code class="lang-auto">import torch
import os
from torch.fx.experimental.proxy_tensor import make_fx
import torch.distributed as dist
import torch.distributed._functional_collectives as funcol
import torch.multiprocessing as mp
from torch._functorch.aot_autograd import aot_module_simplified
from typing import Any, Dict, Iterator, List, Optional, Tuple, Union

def toy_compiler(gm: torch.fx.GraphModule, example_inputs: List[torch.Tensor]):
    print(gm.graph)
    return gm

def toy_backend(gm: torch.fx.GraphModule, example_inputs: List[torch.Tensor]):
    return aot_module_simplified(gm, example_inputs,fw_compiler=toy_compiler)

class MyModule(torch.nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, x, output):
        # torch.distributed.all_reduce(x)   # it can not work!!!
        from torch.distributed.distributed_c10d import _world
        dist.reduce_scatter_tensor(output, x,  group=_world.default_pg)  # it can work
        return x

def example(rank, world_size):
    dist.init_process_group("gloo", rank=rank, world_size=world_size)
    mod = MyModule()
    opt_mod = torch.compile(mod, dynamic=True, fullgraph=True, backend=toy_backend)
    xx = torch.tensor([0, 1, 2, 3], dtype=torch.int32)
    output1 = torch.empty([2], dtype=torch.int32)
    out = opt_mod(xx, output1)
    print(out)


def main():
    world_size = 2
    mp.spawn(example,
             args=(world_size, ),
             nprocs=world_size,
             join=True)


if __name__ == "__main__":
    os.environ["MASTER_ADDR"] = "localhost"
    os.environ["MASTER_PORT"] = "29516"
    main()

</code></pre>
<p>I think the reason why reduce_scatter_tensor works is that it’s in the map</p>
<pre><code class="lang-auto">from torch.distributed.distributed_c10d import (
    all_gather_into_tensor as legacy_allgather,
    reduce_scatter_tensor as legacy_reducescatter,
)

# This dict should contain sets of functions that dynamo is allowed to remap.
# Functions in this set should accept the same args/kwargs 1:1 as their mapping.
traceable_collective_remaps = {
    legacy_allgather: all_gather_tensor_inplace,
    legacy_reducescatter: reduce_scatter_tensor_inplace,
}

</code></pre>
<ol start="2">
<li>In the preceding case, reduce_scatter_tensor must use the group input parameter, which is different from the eager mode. Scripts that will support eager mode in the future can be used in torch.compile? Or we have to make the necessary changes, like</li>
</ol>
<pre><code class="lang-auto">dist.reduce_scatter_tensor(output, x, async_op=False)
------&gt; change to
from torch.distributed.distributed_c10d import _world
dist.reduce_scatter_tensor(output, x,  group=_world.default_pg)
</code></pre>

wconstab: <p>(1)<br>
I think allgather/reducescatter were added to the map precisely because they were needed to unblock the effort of compiling/tracing FSDP code.  It should be possible to also support allreduce.</p>
<p>Caveat: we can’t automatically map/trace the versions of c10d ops where the ‘async=True’ flag is set.  The whole premise of this ‘auto rewrite’ is that it converts an op into our functional equivalent.  It can do this for the synchronous version of the mutating op by calling the functional op first and then copying the output tensor from the functional op back into the storage of the input tensor.  For the async version, dynamo would have to reason about the ‘work object’ that is returned from the collective, and it does not know how to do that.</p>
<p>(2) IIUC your script omits the group argument but that means it uses the default group, which is probably something we could ‘infer’ inside dynamo.</p>
<p>Ill create a task for these changes.</p><aside class="onebox githubissue" data-onebox-src="https://github.com/pytorch/pytorch/issues/112028">
  <header class="source">

      <a href="https://github.com/pytorch/pytorch/issues/112028" target="_blank" rel="noopener nofollow ugc">github.com/pytorch/pytorch</a>
  </header>

  <article class="onebox-body">
    <div class="github-row">
  <div class="github-icon-container" title="Issue">
	  <svg width="60" height="60" class="github-icon" viewBox="0 0 14 16" aria-hidden="true"><path fill-rule="evenodd" d="M7 2.3c3.14 0 5.7 2.56 5.7 5.7s-2.56 5.7-5.7 5.7A5.71 5.71 0 0 1 1.3 8c0-3.14 2.56-5.7 5.7-5.7zM7 1C3.14 1 0 4.14 0 8s3.14 7 7 7 7-3.14 7-7-3.14-7-7-7zm1 3H6v5h2V4zm0 6H6v2h2v-2z"></path></svg>
  </div>

  <div class="github-info-container">
    <h4>
      <a href="https://github.com/pytorch/pytorch/issues/112028" target="_blank" rel="noopener nofollow ugc">dynamo rewrite for allreduce, default group</a>
    </h4>

    <div class="github-info">
      <div class="date">
        opened <span class="discourse-local-date" data-format="ll" data-date="2023-10-25" data-time="16:07:08" data-timezone="UTC">04:07PM - 25 Oct 23 UTC</span>
      </div>


      <div class="user">
        <a href="https://github.com/wconstab" target="_blank" rel="noopener nofollow ugc">
          <img alt="wconstab" src="https://avatars.githubusercontent.com/u/4984825?v=4" class="onebox-avatar-inline" width="20" height="20">
          wconstab
        </a>
      </div>
    </div>

    <div class="labels">
    </div>
  </div>
</div>

  <div class="github-row">
    <p class="github-body-container">### 🚀 The feature, motivation and pitch

motivated by this post
https://discuss<span class="show-more-container"><a href="" rel="noopener" class="show-more">…</a></span><span class="excerpt hidden">.pytorch.org/t/functional-collectives-used-in-torch-compile/190627

TLDR is we currently have dynamo support to rewrite certain forms of non-functional allgather/reducescatter into functional ones, to make the program traceable.

we should probably be able to support the same rewrite for allreduce and possibly also infer the 'default' processgroup in case it is omitted from the arguments to the c10d collective.

cc @voz

### Alternatives

_No response_

### Additional context

_No response_</span></p>
  </div>

  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>