# How to enable GSPMD?

**URL:** <https://discuss.ai.google.dev/t/how-to-enable-gspmd/30524>\
**Category:** General Discussion\
**Tags:** gpu, xla\
**Created:** [October 31, 2022, 9:26am UTC](https://discuss.ai.google.dev/t/how-to-enable-gspmd/30524 "2022-10-31T09:26:51Z")\
**Posts on this page:** 2\
**Page:** 1

<div class="post-metadata">

**Author:** ![Xiaoda\_Zhang](https://d1dlmcr85iqnpo.cloudfront.net/user_avatar/discuss.ai.google.dev/xiaoda_zhang/32/12958_2.png) [@Xiaoda\_Zhang](https://discuss.ai.google.dev/u/Xiaoda_Zhang)\
**Post date:** [October 31, 2022, 9:26am UTC](https://discuss.ai.google.dev/t/how-to-enable-gspmd/30524/1 "2022-10-31T09:26:51Z")

</div>

We are trying to utilize GSPMD ([[2105.04663] GSPMD: General and Scalable Parallelization for ML Computation Graphs](https://arxiv.org/abs/2105.04663)) on a server with 8 GPUS, using the API ‘xla\_sharding.mesh\_split’.

To dump the IRs, we enable the environment variables as: TF\_DUMP\_GRAPH\_PREFIX=/tmp/generated   
TF\_XLA\_FLAGS=“–tf\_xla\_clustering\_debug --tf\_xla\_auto\_jit=2”   
XLA\_FLAGS=“–xla\_dump\_hlo\_as\_text --xla\_dump\_to=/tmp/generated”   
my/tensorflow/program".

But, only 4 IRs are saved: mark\_for\_compilation.pbtxt, mark\_for\_compilation\_annotated.pbtxt, before\_mark\_for\_compilation.pbtxt, before\_increase\_dynamism\_for\_auto\_jit\_pass.pbtxt. None of them is related to SPMD pass.

It seems that our current run does NOT enable GSPMD functionality at all. Is there any tutorial or instructions for us to follow to enable GSPMD on multiple GPUS?

Our testcode is follows (copied and modified from [tensorflow/tensorflow/compiler/xla/experimental/xla\_sharding/xla\_sharding\_test.py at r2.7 · tensorflow/tensorflow · GitHub](https://github.com/tensorflow/tensorflow/blob/r2.7/tensorflow/compiler/xla/experimental/xla_sharding/xla_sharding_test.py)):

from tensorflow.compiler.xla.experimental.xla\_sharding import xla\_sharding  
from tensorflow.python.ops import array\_ops  
from tensorflow.python.ops import math\_ops  
from tensorflow.python.framework import test\_util  
from tensorflow.python.framework import dtypes  
from tensorflow.python.framework import ops  
import numpy as np  
from tensorflow.python.eager import def\_function

class XlaShardingTest(test\_util.TensorFlowTestCase):  
def test\_dot\_split(self):  
@def\_function.function  
def split\_helper(tensor):  
device\_mesh = np.array([[0, 1, 2, 3], [4, 5, 6, 7]])  
split\_tensor = xla\_sharding.mesh\_split(tensor, device\_mesh, [0, 1])  
self.assertIsInstance(split\_tensor, ops.Tensor)  
split\_sharding = xla\_sharding.get\_tensor\_sharding(split\_tensor)  
split\_shape = xla\_sharding.get\_sharding\_tile\_shape(split\_sharding)  
expected\_shape = [2, 4]  
self.assertEqual(expected\_shape, split\_shape)

```
  y_tensor = array_ops.ones([8, 8], dtype=dtypes.float32)
  y_split = xla_sharding.mesh_split(y_tensor, device_mesh, [0, 1])
  result = math_ops.matmul(split_tensor, y_split)
  device_mesh = np.array([[0, 1], [2, 3], [4, 5], [6, 7]])
  result = xla_sharding.mesh_split(result, device_mesh, [0, 1])
  result = math_ops.sqrt(result)
  result = xla_sharding.mesh_split(result, device_mesh, [1, 0])
  return result

in_tensor = 2 * np.sqrt(2) * array_ops.ones([8, 8], dtype=dtypes.float32)
result = split_helper(
    array_ops.ones([8, 8], dtype=dtypes.float32))
self.assertAllEqual(in_tensor, result)

```

if **name** == “ **main** ”:  
xlasharding = XlaShardingTest()  
xlasharding.test\_dot\_split()

---

<div class="post-metadata">

**Author:** ![Surya\_Y](https://d1dlmcr85iqnpo.cloudfront.net/user_avatar/discuss.ai.google.dev/surya_y/32/12926_2.png) [@Surya\_Y](https://discuss.ai.google.dev/u/Surya_Y)\
**Post date:** [November 2, 2022, 8:55am UTC](https://discuss.ai.google.dev/t/how-to-enable-gspmd/30524/3 "2022-11-02T08:55:06Z")

</div>

Hi @Xiaoda_Zhang ,the test code works fine as per [gist](https://colab.research.google.com/gist/SuryanarayanaY/f085e5f4040bac0b5009bad95cec5dfb/58363.ipynb) but not sure about the GSPMD enabling things.Anyway can you go through the content of the generated files from code execution and share if any useful information there.  
Could you please confirm details like how you have enabled the GPUs and which distribution strategy followed etc to dig into the problem.