Export module to StableHLO with communication collectives

username: hsnbrg
cooked: <p>Hi,</p>
<p>I’m trying to export a module to StableHLO with the sharding specification and also the communication collectives inserted.<br>
I’m using <code>torch_xla</code> 2.2.<br>
This is what my module looks like</p>
<pre><code class="lang-auto">
import torch
import  torch.nn as nn
import torch_xla
torch_xla.runtime.use_spmd()

class ScaleModule(nn.Module):
    def __init__(self, scale):
        super().__init__()
        self.scale = scale

    def forward(self, input: torch.Tensor):
        result = self.scale * input
        return result
</code></pre>
<p>And I’m performing the following steps</p>
<pre><code class="lang-auto">from torch_xla.distributed.spmd import xla_sharding as xs
from torch_xla.distributed.spmd import Mesh
from torch.export import export
from torch_xla.stablehlo import exported_program_to_stablehlo
import torch_xla.core.xla_model as xm

input_tensor = torch.randn([8, 1, 8, 16])
mesh = Mesh((0, 1, 2, 3, 4, 5, 6, 7), (2, 4), ('x', 'y'))
partition_spec = ('x', None, None, None)
sharded_tensor = xs.mark_sharding(input_tensor.to(xm.xla_device()), mesh, partition_spec)
exported = export(model.to(xm.xla_device()), args=(sharded_tensor.global_tensor, ))
stable = exported_program_to_stablehlo(exported)
</code></pre>
<p>I do see the sharding information on the tensor as part of the <code>mhlo.sharding</code> attribute in the generated StableHLO, however I don’t see the module/operations getting sharded based on this information. From the user guide, it looks like the <code>XLA SpmdPartitioner</code> pass is responsible for this, however I’m not sure what invokes this pass.</p>
<p>I’d appreciate any help. Thanks!</p>

username: Yeounoh_Chung
cooked: <p>The stablehlo you generate still just captures free compilation IR with sharding annotations. The SPMD partitioning pass happens during runtime, as part of the XLA compilation pass.</p>
<p>Try printing/accessing the output of your exported program, with this environment variable set: <code>XLA_FLAGS="--xla_dump_to={dir_path_to_your_hlo_dump}"</code> You should be able to post optimization pass HLO, which should contain the results of SPMD partitinoing.</p>
<p>Hope this helps.</p>

username: hsnbrg
cooked: <p>I see the following StableHLO.</p>
<pre><code class="lang-auto">module @IrToHlo.6 attributes {mhlo.cross_program_prefetches = [], mhlo.is_dynamic = false, mhlo.use_auto_spmd_partitioning = false} {
  func.func @main(%arg0: tensor&lt;f32&gt; {mhlo.sharding = "{replicated}"}, %arg1: tensor&lt;8x1x8x16xf32&gt; {mhlo.sharding = "{devices=[2,1,1,1,4]0,1,2,3,4,5,6,7 last_tile_dim_replicate}"}) -&gt; tensor&lt;8x1x8x16xf32&gt; {
    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor&lt;f32&gt;) -&gt; tensor&lt;8x1x8x16xf32&gt;
    %1 = stablehlo.multiply %arg1, %0 : tensor&lt;8x1x8x16xf32&gt;
    return %1 : tensor&lt;8x1x8x16xf32&gt;
  }
}
</code></pre>
<p>I’m not sure I follow this sentence</p>
<blockquote>
<p>You should be able to post optimization pass HLO, which should contain the results of SPMD partitinoing.</p>
</blockquote>
<p>I do realize that I need to run some pass to obtain an IR post SPMD partitioning. Although I’m not sure how to do that.</p>

username: Yeounoh_Chung
cooked: <aside class="quote no-group" data-username="Yeounoh_Chung" data-post="2" data-topic="199583">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/yeounoh_chung/48/63164_2.png" class="avatar"> Yeounoh_Chung:</div>
<blockquote>
<p>Try printing/accessing the output of your exported program, with this environment variable set: <code>XLA_FLAGS="--xla_dump_to={dir_path_to_your_hlo_dump}"</code> You should be able to post optimization pass HLO, which should contain the results of SPMD partitinoing.</p>
</blockquote>
</aside>
<p>The StableHLO is pre-optimization pass in XLA.  PyTorch/XLA takes it and converts to HLO and then compile to an XLA executable. During the compilation we run SPMD partitioning. So I was suggesting that:</p>
<ol>
<li>your StableHLO won’t show the sharding results,</li>
<li>you should dump XLA post-optimization HLO to see the sharding results.</li>
</ol>

username: hsnbrg
cooked: <p>By compilation, I assume you mean running <code>torch.compile</code> on the exported program that contains the sharding annotations on the tensor?</p>
<p>In the post optimization file, I see the following</p>
<pre><code class="lang-auto">HloModule SyncTensorsGraph.3, entry_computation_layout={(f32[8,1,8,16]{3,2,1,0})-&gt;(f32[8,1,8,16]{3,2,1,0})}, allow_spmd_sharding_propagation_to_output={true}

ENTRY SyncTensorsGraph.3 {
  p0.1 = f32[8,1,8,16]{3,2,1,0} parameter(0), sharding={replicated}
  copy = f32[8,1,8,16]{3,2,1,0} copy(p0.1)
  ROOT tuple = (f32[8,1,8,16]{3,2,1,0}) tuple(copy)
}
</code></pre>
<p>Is there a way to convert the above IR to StableHLO?<br>
To give you some context as to why I’m asking this, I’m trying to use XLA for export to StableHLO alone. I have my own compilation flow that consumes StableHLO and compiles it for a custom backend.<br>
And I was wondering if there was a way to have StableHLO represent a module post SPMD so that my custom compilation flow can consume it.</p>
<p>I appreciate the help! <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=12" title=":slight_smile:" class="emoji" alt=":slight_smile:" loading="lazy" width="20" height="20"></p>

username: Yeounoh_Chung
cooked: <p>Hi <a class="mention" href="/u/hsnbrg">@hsnbrg</a> ,</p>
<aside class="quote no-group" data-username="Yeounoh_Chung" data-post="4" data-topic="199583">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/yeounoh_chung/48/63164_2.png" class="avatar"> Yeounoh_Chung:</div>
<blockquote>
<p>The StableHLO is pre-optimization pass in XLA. PyTorch/XLA takes it and converts to HLO and then compile to an XLA executable.</p>
</blockquote>
</aside>
<p><code>torch.compile</code> is what gets you to the IR. The XLA compilation to XLA executable happens later via PyTorch/XLA (<code>torch_xla</code>). “The above IR” is the HLO. And to you additional questions,</p>
<ul>
<li>
<p>yes, you can try convert HLO to MHLO and then to StableHLO. That should work.</p>
</li>
<li>
<p>the StableHLO you get via export, won’t capture the post XLA optimization result. I think you should go through XLA compilation to get one. This may be a bit hacky for your workflow… unfortunately, we don’t have an API for that.</p>
</li>
</ul>

username: hsnbrg
cooked: <p>Just to make sure I’m understanding this correctly, <code>torch.compile</code> gives me HLO (hopefully respecting the sharding annotations on the exported program that I pass to it).<br>
However, there is another compilation flow via PyTorch/XLA (<code>torch_xla</code>) that potentially consumes HLO? If so, how might I use this compilation flow?</p>
<p>Again, thank you!</p>

username: hsnbrg
cooked: <p>Also, FWIW, I get the same HLO using <code>torch.compile</code> on exported programs pertaining to the sharded vs the unsharded case.<br>
Is there no way to make <code>torch.compile</code> obey the sharding annotations even if the underlying device doesn’t conform to it?</p>

username: Yeounoh_Chung
cooked: <p>Let me summarize some key points…</p>
<ul>
<li><code>torch.compile</code> is not XLA compilation – but <code>torch.compile</code> with <code>openxla</code> backend will call XLA compilation eventually.</li>
<li>Having said that, “The StableHLO is pre-optimization pass in XLA. PyTorch/XLA takes it and converts to HLO and then compile to an XLA executable.” the compilation here is XLA compilation, doesn’t involve <code>torch.compile</code>.</li>
<li>you won’t get the HLO post-optimization, as I explained in the above comments. Currently, the only way is for you to use <code>--xla_dump_to</code>.</li>
</ul>
<aside class="quote no-group" data-username="hsnbrg" data-post="8" data-topic="199583">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/hsnbrg/48/61753_2.png" class="avatar"> hsnbrg:</div>
<blockquote>
<p>Is there no way to make <code>torch.compile</code> obey the sharding annotations even if the underlying device doesn’t conform to it?</p>
</blockquote>
</aside>
<ul>
<li>currently no. XLA compilation is also jit, so happens as part of XLA execution, and must confirm to the current platform.</li>
<li>in the future, yes. We are working on an AOT compilation service in PyTorch/XLA.</li>
</ul>

username: hsnbrg
cooked: <p>Noted. Thank you! I’m fiddling around with <code>hlo-opt</code> to learn more about what you just described.</p>