Add support for in-place ops with self tensors in dynamo bridge

[wonjoolee95]

Add support for in-place ops with self tensors in dynamo bridge

In models where there there is an in-place op on self.tensor, the self.tensor is not part of the xla_args. To show an example, consider the following model:

import torch
import torch.nn as nn

import torch_xla.core.xla_model as xm

device = xm.xla_device()
print(f'deivce={device}')

class TestModel(nn.Module):
    def __init__(self, device=None):
        super().__init__()
        self.b = torch.zeros((3, 3), device=device)

    def forward(self, indexes, buffer_update, input):
        self.b.index_copy_(0, indexes, buffer_update)
        output = input + self.b
        return output

Currently, running this model with dynamo will error out with a Check failed: HasValue() error:

RuntimeError: ./torch_xla/csrc/runtime/pjrt_computation_client.h:166 : Check failed: HasValue() 

And looking at xla_args in our xla/torch_xla/core/dynamo_bridge.py at function extract_compiled_graph, we can see that xla_args passed from dynamo is empty:

[WONJOO | dynamo_bridge.py] xla_args=()

And as a result, the self.tensor is not materialized throughout the dynamo code patch, eventually causing the Check failed: HasValue() later when it's accessed.

As for a fix, we manually include self.tensor as part of the xla_args by calling the xla_model.named_buffers() that returns the the self.tensor as such:

Code:
for name, buffer in xla_model.named_buffers():
  print(f'[WONJOO | dynamo_bridge.py] name={name}')
  print(f'[WONJOO | dynamo_bridge.py] buffer={buffer}')

Output:
[WONJOO | dynamo_bridge.py] name=L__self___b
[WONJOO | dynamo_bridge.py] buffer=tensor([[0., 0., 0.],
        [0., 0., 0.],
        [0., 0., 0.]], device='xla:0')

[wonjoolee95]

With the changes now, I can see that the model in the PR description with dynamo is now passing:

deivce=xla:0
-----CPU-----
before: cpu_model.b=tensor([[0., 0., 0.],
        [0., 0., 0.],
        [0., 0., 0.]])
output=tensor([[0.9724, 1.4463, 1.0837],
        [0.4915, 0.1398, 0.0953],
        [0.1789, 0.6370, 0.4601]])
after: cpu_model.b=tensor([[0.2868, 0.7033, 0.8447],
        [0.0000, 0.0000, 0.0000],
        [0.0000, 0.0000, 0.0000]])
-----XLA-----
before: model.b=tensor([[0., 0., 0.],
        [0., 0., 0.],
        [0., 0., 0.]], device='xla:0')
output_xla=tensor([[0.9724, 1.4463, 1.0837],
        [0.4915, 0.1398, 0.0953],
        [0.1789, 0.6370, 0.4601]], device='xla:0')
after: model.b=tensor([[0.2868, 0.7033, 0.8447],
        [0.0000, 0.0000, 0.0000],
        [0.0000, 0.0000, 0.0000]], device='xla:0')

Also added a unit test for in-place ops.

[wonjoolee95]

Need a way to explicitly materialize the tensor, hence the _xla_sync_multi call.

[JackCaoG]

why do we need this? Didn't we do mark_step on the entrance of the extract graph?

[wonjoolee95]

The allows_single_node_partition=True flag ensures that the partitioned module (even when it is one single large partition) goes into our own extract_internal call.

[wonjoolee95]

Some metrics failures with Dynamo, most likely due to the new _xla_sync_multi call. Let me make sure that only the metric tests are failing and that the other correctness tests are still passing.

[wonjoolee95]

Some metrics failures with Dynamo, most likely due to the new _xla_sync_multi call. Let me make sure that only the metric tests are failing and that the other correctness tests are still passing.

Surrounded the _xla_sync_multi in an if-statement so it only runs if there is something in the named_buffers (i.e. there is a self.tensor). So the metrics shouldn't be impacted anymore, I'll let CI verify.

[wonjoolee95]

Seems like the allows_single_node_partition=True affected the CompileTime for CpuFallback related tests only (became one less). While I make sense of it, I updated the tests for now to let the CI continue to run.

[wonjoolee95]

Hmm, putting some debugging lines into the xm.mark_step() call (https://github.com/pytorch/xla/blob/master/torch_xla/core/dynamo_bridge.py#L392), I can see that the SyncLiveTensorsGraph actually does seem to materialize all the tensors, including the self.tensor. The following prints have been added at https://github.com/pytorch/xla/blob/master/torch_xla/csrc/xla_graph_executor.cpp#L374, for-looping through the tensors:

[WONJOO | xla_graph_executor.cpp] tensors.size()=4
[WONJOO | xla_graph_executor.cpp] tensor->GetUniqueId()=1
[WONJOO | xla_graph_executor.cpp] tensor->shape().get().ToString()=f32[3]
[WONJOO | xla_graph_executor.cpp] tensor->GetUniqueId()=2
[WONJOO | xla_graph_executor.cpp] tensor->shape().get().ToString()=f32[3,3]
[WONJOO | xla_graph_executor.cpp] tensor->GetUniqueId()=3
[WONJOO | xla_graph_executor.cpp] tensor->shape().get().ToString()=s64[3]
[WONJOO | xla_graph_executor.cpp] tensor->GetUniqueId()=5
[WONJOO | xla_graph_executor.cpp] tensor->shape().get().ToString()=f32[5,3]

This is logs from running the newly added unit test DynamoInferenceBasicTest.test_simple_model_with_in_place_ops. And as we can see, the f32[5,3] with id 5 is the self.tensor and it does seem to get materialized.

However, with just this mark_step() and even with setting the wait=True flag, it still seems to fail with the buffer with shape f32[5,3] on device CPU:0 is null error if we remove that explicit xla_sync_multi call.

[wonjoolee95]

Okay, seems like the comment at #5309 (comment) may not be true. In the initial mark_step() call, I can see the following logs:

[WONJOO | init_python_bindings.cpp] at StepMarker
[WONJOO | xla_graph_executor.cpp] at SyncLiveTensorsGraph
[WONJOO | xla_graph_executor.cpp] tensors.size()=4
[WONJOO | xla_graph_executor.cpp] tensor->GetUniqueId()=1
[WONJOO | xla_graph_executor.cpp] tensor->shape().get().ToString()=f32[3]
[WONJOO | xla_graph_executor.cpp] tensor->GetUniqueId()=2
[WONJOO | xla_graph_executor.cpp] tensor->shape().get().ToString()=f32[3,3]
[WONJOO | xla_graph_executor.cpp] tensor->GetUniqueId()=3
[WONJOO | xla_graph_executor.cpp] tensor->shape().get().ToString()=s64[3]
[WONJOO | xla_graph_executor.cpp] tensor->GetUniqueId()=5
[WONJOO | xla_graph_executor.cpp] tensor->shape().get().ToString()=f32[5,3]

I thought f32[5,3] was the self.tensor. However, when I just try to add back the `_xla_sync_multi(self_tensors) call, I can see:

[WONJOO | init_python_bindings.cpp] SyncTensors, wait=0, sync_xla_data=1, warm_up_cache_only=0
[WONJOO | xla_graph_executor.cpp] at SyncTensorsGraph
[WONJOO | xla_graph_executor.cpp] wait=0
[WONJOO | xla_graph_executor.cpp] sync_ltc_data=1
[WONJOO | xla_graph_executor.cpp] warm_up_cache_only=0
[WONJOO | xla_graph_executor.cpp] tensors->size()=1
[WONJOO | xla_graph_executor.cpp] tensor->GetUniqueId()=12
[WONJOO | xla_graph_executor.cpp] tensor->shape().get().ToString()=f32[5,3]

It actually looks like the self_tensor is has GetUniqueId()=12. And this tensor with GetUniqueId()=12 is not included in the part of the initial mark_step call.

[JackCaoG]

Hmm, let me try to repo

[JackCaoG]

I can repo, I added a debug message

  for name, buffer in xla_model.named_buffers():
    if "self" in name:
      self_tensors.append(buffer)
      print(torch_xla._XLAC._get_xla_tensor_debug_info(buffer))

and in the

@@ -87,7 +89,9 @@ std::vector<XLATensorPtr> XLAGraphExecutor::DeviceContextArena::GetLiveTensors(
       auto data =
           std::dynamic_pointer_cast<XLATensor::Data>(uid_wptr.second.lock());
       if (data != nullptr) {
-        tensors.push_back(XLATensor::Create(std::move(data)));
+        auto t = XLATensor::Create(std::move(data));
+        std::cerr << "add tensor with id " << t->GetUniqueId() << "\n";
+        tensors.push_back(t);
       }

I saw

add tensor with id 1
add tensor with id 2
add tensor with id 3
add tensor with id 5
XLATensor {
TensorID: 11
Device: CPU:0
XLA Shape: f32[5,3]
IR: [] aten::index_add_, xla_shape=f32[5,3]{1,0}
XLAData: None
Tensor on host: None
}

add tensor with id 1
add tensor with id 2
add tensor with id 3
add tensor with id 6
add tensor with id 7
add tensor with id 8
add tensor with id 10
add tensor with id 11

What this means is that we do include the self.x in the mark_step somewhere. Then the next step is for me to figure out which mark_step. I added a pdb to mark_step call

--- a/torch_xla/core/xla_model.py
+++ b/torch_xla/core/xla_model.py
@@ -797,6 +797,7 @@ def _run_step_closures():
 
 
 def mark_step(wait=False):
+  import pdb; pdb.set_trace()
   if xu.getenv_as('XLA_EMIT_STEPLOG', bool, False):
     print(
         'torch_xla.core.xla_model::mark_step\n',

and found that first mark_step is in

391  	def extract_compiled_graph(xla_model, xla_args):
392  	  # This call is critical to make sure xla_args' tensor id show up in graph_input_tensor_ids
393  ->	  xm.mark_step()

which result in

add tensor with id 1
add tensor with id 2
add tensor with id 3
add tensor with id 5

Then I didn't see the second makr_step before seeing

add tensor with id 1
add tensor with id 2
add tensor with id 3
add tensor with id 6
add tensor with id 7
add tensor with id 8
add tensor with id 10
add tensor with id 11

which suggested that second execution is not triggered by a mark_step. If I just keep stepping the python function I can see it is actually _clear_pending_ir triggers the clean up. so

> /src/pytorch/xla/torch_xla/core/dynamo_bridge.py(444)extract_compiled_graph()
-> torch_xla._XLAC._clear_pending_irs(str(xm.xla_device()))
(Pdb) 
add tensor with id 1
add tensor with id 2
add tensor with id 3
add tensor with id 6
add tensor with id 7
add tensor with id 8
add tensor with id 10
add tensor with id 12

This make sense since if you look at how _clear_pending_irs works, it pretty much just collect all tensors and remove pending ir. Now the question goes back to why _clear_pending_irs will clean the self.x buffer

[JackCaoG]

The answer to above question is clear, because self.self_tensor actually has a IR, which is IR: [] aten::index_add_, xla_shape=f32[5,3]{1,0}. I think this is caused by tracing the model.

I think the issue is FallBackNodeCollector and CapabilityBasedPartitioner perform the tracing on self.self_tensor. For all of the xla_args we avoid this issue by

  cloned_xla_args = [
      torch.clone(xla_arg) if isinstance(xla_arg, torch.Tensor) else xla_arg
      for xla_arg in xla_args
  ]

and replace them after tracing. We need to do the same thing for the named_buffer

[JackCaoG]

@wonjoolee95 Is this ready for review?

[wonjoolee95]

Execution numbers for fallback-related tests only changed because of the allows_single_node_partition=True mentioned below.

[wonjoolee95]

@wonjoolee95 Is this ready for review?

Yep, should be ready to review now. All dynamo unit tests passing locally:

(base) wonjoo@wonjoo-cpu:~/pytorch/xla$ python test/dynamo/test_dynamo.py
/home/wonjoo/miniconda3/lib/python3.10/site-packages/torchvision/io/image.py:13: UserWarning: Failed to load image Python extension: 'libc10_cuda.so: cannot open shared object file: No such file or directory'If you don't plan on using image functionality from `torchvision.io`, you can ignore this warning. Otherwise, there might be something wrong with your environment. Did you have `libjpeg` or `libpng` installed before building `torchvision` from source?
  warn(
...........
----------------------------------------------------------------------
Ran 11 tests in 28.067s

OK

[wonjoolee95]

Hmm, the failing CI seems to succeed locally:

(base) wonjoo@wonjoo-cpu:~/pytorch/xla$ python test/dynamo/test_dynamo.py DynamoInferenceBasicTest.test_simple_model_with_in_place_ops
/home/wonjoo/miniconda3/lib/python3.10/site-packages/torchvision/io/image.py:13: UserWarning: Failed to load image Python extension: 'libc10_cuda.so: cannot open shared object file: No such file or directory'If you don't plan on using image functionality from `torchvision.io`, you can ignore this warning. Otherwise, there might be something wrong with your environment. Did you have `libjpeg` or `libpng` installed before building `torchvision` from source?
  warn(
.
----------------------------------------------------------------------
Ran 1 test in 0.112s

OK

Looking into it.

[JackCaoG]

seems like cpu test still failed?

[wonjoolee95]

Hmm, the behavior of the newly added test test_simple_model_with_in_place_ops is odd. Trying to locally reproduce, I can see it succeeds. But if I keep running it consecutively, I'll eventually see a fail (5 consecutive runs will result in 4 successes and 1 failure):

(base) wonjoo@wonjoo-cpu:~/pytorch/xla$ python test/dynamo/test_dynamo.py 
.
----------------------------------------------------------------------
Ran 1 test in 0.120s

OK
(base) wonjoo@wonjoo-cpu:~/pytorch/xla$ python test/dynamo/test_dynamo.py 
.
----------------------------------------------------------------------
Ran 1 test in 0.120s

OK
(base) wonjoo@wonjoo-cpu:~/pytorch/xla$ python test/dynamo/test_dynamo.py 
(base) wonjoo@wonjoo-cpu:~/pytorch/xla$ python test/dynamo/test_dynamo.py 
.
----------------------------------------------------------------------
Ran 1 test in 0.120s

OK
(base) wonjoo@wonjoo-cpu:~/pytorch/xla$ python test/dynamo/test_dynamo.py DynamoInferenceBasicTest.test_simple_model_with_in_place_ops
F
======================================================================
FAIL: test_simple_model_with_in_place_ops (__main__.DynamoInferenceBasicTest)
----------------------------------------------------------------------
Traceback (most recent call last):
  File "/home/wonjoo/pytorch/xla/test/dynamo/test_dynamo.py", line 112, in test_simple_model_with_in_place_ops
    self.assertTrue(torch.allclose(res_cpu, res_xla_dynamo.cpu()))
AssertionError: False is not true

----------------------------------------------------------------------
Ran 1 test in 0.122s

FAILED (failures=1)
(base) wonjoo@wonjoo-cpu:~/pytorch/xla$ 

[JackCaoG]

can you print the res_cpu and res_xla_dynamo?

[wonjoolee95]

So it seems like sometimes that res_xla_dynamo isn't updated properly:

(base) wonjoo@wonjoo-cpu:~/pytorch/xla$ python test/dynamo/test_dynamo.py DynamoInferenceBasicTest.test_simple_model_with_in_place_ops
[WONJOO] res_cpu=tensor([[ 2.,  3.,  4.],
        [ 1.,  1.,  1.],
        [ 8.,  9., 10.],
        [ 1.,  1.,  1.],
        [ 5.,  6.,  7.]])
[WONJOO] res_xla_dynamo.cpu()=tensor([[ 2.,  3.,  4.],
        [ 1.,  1.,  1.],
        [ 8.,  9., 10.],
        [ 1.,  1.,  1.],
        [ 5.,  6.,  7.]])
.
----------------------------------------------------------------------
Ran 1 test in 0.121s

OK
(base) wonjoo@wonjoo-cpu:~/pytorch/xla$ python test/dynamo/test_dynamo.py DynamoInferenceBasicTest.test_simple_model_with_in_place_ops
[WONJOO] res_cpu=tensor([[ 2.,  3.,  4.],
        [ 1.,  1.,  1.],
        [ 8.,  9., 10.],
        [ 1.,  1.,  1.],
        [ 5.,  6.,  7.]])
[WONJOO] res_xla_dynamo.cpu()=tensor([[ 2.,  3.,  4.],
        [ 1.,  1.,  1.],
        [ 8.,  9., 10.],
        [ 1.,  1.,  1.],
        [ 5.,  6.,  7.]])
.
----------------------------------------------------------------------
Ran 1 test in 0.120s

OK
(base) wonjoo@wonjoo-cpu:~/pytorch/xla$ python test/dynamo/test_dynamo.py DynamoInferenceBasicTest.test_simple_model_with_in_place_ops
[WONJOO] res_cpu=tensor([[ 2.,  3.,  4.],
        [ 1.,  1.,  1.],
        [ 8.,  9., 10.],
        [ 1.,  1.,  1.],
        [ 5.,  6.,  7.]])
[WONJOO] res_xla_dynamo.cpu()=tensor([[1., 1., 1.],
        [1., 1., 1.],
        [1., 1., 1.],
        [1., 1., 1.],
        [1., 1., 1.]])
F

Thinking how this fails only once in a while, I thought it may be related to the mark_step(), so tried updating all the calls to mark_step(wait=True) but still seeing the occasional failure.

[JackCaoG]

Let me take a look

[JackCaoG]

I am able to repo the random failure, looking into it

[JackCaoG]

I noticed something weird, if I tried to dump the info about xla_arg and self_arg in the beginning of the extract internal with

def extract_internal(xla_model: torch.fx.GraphModule, self_args):
  xla_args = xla_model.xla_args
  print('xla arg\n')
  for xla_arg in xla_args:
    print(torch_xla._XLAC._get_xla_tensor_debug_info(xla_arg))
  print('self arg\n')
  for self_arg in self_args:
    print(torch_xla._XLAC._get_xla_tensor_debug_info(self_arg))  

xla arg

XLATensor {
TensorID: 1
Device: CPU:0
XLA Shape: f32[3]
ShardingSpec: None
IR: None
XLAData: 
  Data Device: CPU:0
  Data Shape: f32[3]
  Data Handle: 140350739976128
Tensor on host: None
}

XLATensor {
TensorID: 13
Device: CPU:0
XLA Shape: f32[5,3]
ShardingSpec: None
IR: None
XLAData: 
  Data Device: CPU:0
  Data Shape: f32[5,3]
  Data Handle: 140350739978224
Tensor on host: None
}

XLATensor {
TensorID: 3
Device: CPU:0
XLA Shape: s64[3]
ShardingSpec: None
IR: None
XLAData: 
  Data Device: CPU:0
  Data Shape: s64[3]
  Data Handle: 140350739977440
Tensor on host: None
}

XLATensor {
TensorID: 2
Device: CPU:0
XLA Shape: f32[3,3]
ShardingSpec: None
IR: None
XLAData: 
  Data Device: CPU:0
  Data Shape: f32[3,3]
  Data Handle: 140350739976688
Tensor on host: None
}

self arg

XLATensor {
TensorID: 13
Device: CPU:0
XLA Shape: f32[5,3]
ShardingSpec: None
IR: None
XLAData: 
  Data Device: CPU:0
  Data Shape: f32[5,3]
  Data Handle: 140350739978224
Tensor on host: None
}

Note that Tensor with tensor ID 13 showed up both in xla_arg and self_arg. This kind of suggest we don't need to special casing the self_arg.. Need to look into how much of that is caused by the change in this pr.

[JackCaoG]

ok I am a bit confuse by the additional layer of extract_compiled_graph and extract_internal, and we need to do this arg clone twice is confusing.. need to spend some time reading it again.

[JackCaoG]

hmm I am confuse, if I print in the extract_compiled_graph, I see

(Pdb) print(len(self_args))
1
(Pdb) print(len(xla_args))
3

if I look at extract_internal

(Pdb) len(self_args)
1
(Pdb) len(xla_args)
4

for some reason after the partitioner, self_arg becomes part of the xla_arg(or xla_model.xla_args)

[JackCaoG]

I dump the graph by adding

  torch_xla._XLAC._xla_warm_up_cache(args_and_out, [])
 + print(torch_xla._XLAC._get_xla_tensors_hlo(args_and_out))

and for the passing one I see

HloModule IrToHlo.18, entry_computation_layout={(f32[3,3]{1,0}, s64[3]{0}, f32[5,3]{1,0}, f32[3]{0})->(f32[5,3]{1,0}, f32[5,3]{1,0})}

%ScatterCombiner.4 (p0.5: f32[], p1.6: f32[]) -> f32[] {
  %p0.5 = f32[] parameter(0)
  ROOT %p1.6 = f32[] parameter(1)
}

ENTRY %IrToHlo.18 (p0.1: f32[3,3], p1.2: s64[3], p2.3: f32[5,3], p3.14: f32[3]) -> (f32[5,3], f32[5,3]) {
  %p2.3 = f32[5,3]{1,0} parameter(2)
  %p1.2 = s64[3]{0} parameter(1)
  %p0.1 = f32[3,3]{1,0} parameter(0)
  %scatter.7 = f32[5,3]{1,0} scatter(f32[5,3]{1,0} %p2.3, s64[3]{0} %p1.2, f32[3,3]{1,0} %p0.1), update_window_dims={1}, inserted_window_dims={0}, scatter_dims_to_operand_dims={0}, index_vector_dim=1, to_apply=%ScatterCombiner.4
  %p3.14 = f32[3]{0} parameter(3)
  %broadcast.15 = f32[5,3]{1,0} broadcast(f32[3]{0} %p3.14), dimensions={1}
  %constant.8 = f32[] constant(1)
  %reshape.9 = f32[1,1]{1,0} reshape(f32[] %constant.8)
  %broadcast.10 = f32[1,1]{1,0} broadcast(f32[1,1]{1,0} %reshape.9), dimensions={0,1}
  %reshape.11 = f32[] reshape(f32[1,1]{1,0} %broadcast.10)
  %broadcast.12 = f32[5,3]{1,0} broadcast(f32[] %reshape.11), dimensions={}
  %multiply.13 = f32[5,3]{1,0} multiply(f32[5,3]{1,0} %scatter.7, f32[5,3]{1,0} %broadcast.12)
  %add.16 = f32[5,3]{1,0} add(f32[5,3]{1,0} %broadcast.15, f32[5,3]{1,0} %multiply.13)
  ROOT %tuple.17 = (f32[5,3]{1,0}, f32[5,3]{1,0}) tuple(f32[5,3]{1,0} %scatter.7, f32[5,3]{1,0} %add.16)
}

and for the failing one I see

HloModule IrToHlo.18, entry_computation_layout={(f32[3,3]{1,0}, s64[3]{0}, f32[5,3]{1,0}, f32[3]{0})->(f32[5,3]{1,0}, f32[5,3]{1,0})}

%ScatterCombiner.4 (p0.5: f32[], p1.6: f32[]) -> f32[] {
  %p0.5 = f32[] parameter(0)
  ROOT %p1.6 = f32[] parameter(1)
}

ENTRY %IrToHlo.18 (p0.1: f32[3,3], p1.2: s64[3], p2.3: f32[5,3], p3.14: f32[3]) -> (f32[5,3], f32[5,3]) {
  %p2.3 = f32[5,3]{1,0} parameter(2)
  %p1.2 = s64[3]{0} parameter(1)
  %p0.1 = f32[3,3]{1,0} parameter(0)
  %scatter.7 = f32[5,3]{1,0} scatter(f32[5,3]{1,0} %p2.3, s64[3]{0} %p1.2, f32[3,3]{1,0} %p0.1), update_window_dims={1}, inserted_window_dims={0}, scatter_dims_to_operand_dims={0}, index_vector_dim=1, to_apply=%ScatterCombiner.4
  %p3.14 = f32[3]{0} parameter(3)
  %broadcast.15 = f32[5,3]{1,0} broadcast(f32[3]{0} %p3.14), dimensions={1}
  %constant.8 = f32[] constant(1)
  %reshape.9 = f32[1,1]{1,0} reshape(f32[] %constant.8)
  %broadcast.10 = f32[1,1]{1,0} broadcast(f32[1,1]{1,0} %reshape.9), dimensions={0,1}
  %reshape.11 = f32[] reshape(f32[1,1]{1,0} %broadcast.10)
  %broadcast.12 = f32[5,3]{1,0} broadcast(f32[] %reshape.11), dimensions={}
  %multiply.13 = f32[5,3]{1,0} multiply(f32[5,3]{1,0} %p2.3, f32[5,3]{1,0} %broadcast.12)
  %add.16 = f32[5,3]{1,0} add(f32[5,3]{1,0} %broadcast.15, f32[5,3]{1,0} %multiply.13)
  ROOT %tuple.17 = (f32[5,3]{1,0}, f32[5,3]{1,0}) tuple(f32[5,3]{1,0} %scatter.7, f32[5,3]{1,0} %add.16)
}

the difference is in this line, the one produce the correct output uses

  %multiply.13 = f32[5,3]{1,0} multiply(f32[5,3]{1,0} %scatter.7, f32[5,3]{1,0} %broadcast.12)

the one gives incorrect output does

  %multiply.13 = f32[5,3]{1,0} multiply(f32[5,3]{1,0} %p2.3, f32[5,3]{1,0} %broadcast.12)

scatter is pretty much just the result of the self.self_tensor.index_copy_(0, index, copy_tensor). To conclude, the difference in result depends on whether add happens on the result of the inde_copy or the origional self.self_tensor. In either case, we can see that scatter.7 is presented in the HLO, which means self.self_tensor should be updated.

I can confirm that by

-> self.assertTrue(torch.allclose(res_cpu, res_xla_dynamo.cpu()))
(Pdb) compiled_model.self_tensor
tensor([[1., 2., 3.],
        [0., 0., 0.],
        [7., 8., 9.],
        [0., 0., 0.],
        [4., 5., 6.]], device='xla:0')
(Pdb) cpu_model.self_tensor
tensor([[1., 2., 3.],
        [0., 0., 0.],
        [7., 8., 9.],
        [0., 0., 0.],
        [4., 5., 6.]])

[JackCaoG]

Ok I think I found the problem, it is actually not part of the pytorch/xla. I enabled the TORCH_XLA_DEBUG=1 which will dump the fx graph, in the passing case I see

def forward(self, l__self___self_tensor, l_index_ : torch.Tensor, l_copy_tensor_ : torch.Tensor, l_input_tensor_ : torch.Tensor):
    index_copy_ = l__self___self_tensor.index_copy_(0, l_index_, l_copy_tensor_);  l_index_ = l_copy_tensor_ = None
    add = l_input_tensor_ + l__self___self_tensor;  l_input_tensor_ = l__self___self_tensor = None
    return add

in the failing case I see

def forward(self, l_input_tensor_ : torch.Tensor, l__self___self_tensor, l_index_ : torch.Tensor, l_copy_tensor_ : torch.Tensor):
    add = l_input_tensor_ + l__self___self_tensor;  l_input_tensor_ = None
    index_copy_ = l__self___self_tensor.index_copy_(0, l_index_, l_copy_tensor_);  l__self___self_tensor = l_index_ = l_copy_tensor_ = None
    return add

In the failing, add happens before index_copy. Trying to figure out whether this is related to partitioner.

[JackCaoG]

If I dump the fx graph before the partitioner , I see

def forward(self, L_index_ : torch.Tensor, L_copy_tensor_ : torch.Tensor, L_input_tensor_ : torch.Tensor):
    l_index_ = L_index_
    l_copy_tensor_ = L_copy_tensor_
    l_input_tensor_ = L_input_tensor_
    l__self___self_tensor = self.L__self___self_tensor
    index_copy_ = l__self___self_tensor.index_copy_(0, l_index_, l_copy_tensor_);  l_index_ = l_copy_tensor_ = None
    add = l_input_tensor_ + l__self___self_tensor;  l_input_tensor_ = l__self___self_tensor = None
    return (add,)

after the partitioner

def forward(self, l_input_tensor_ : torch.Tensor, l__self___self_tensor, l_index_ : torch.Tensor, l_copy_tensor_ : torch.Tensor):
    add = l_input_tensor_ + l__self___self_tensor;  l_input_tensor_ = None
    index_copy_ = l__self___self_tensor.index_copy_(0, l_index_, l_copy_tensor_);  l__self___self_tensor = l_index_ = l_copy_tensor_ = None
    return add

so the issue is in the partitioner messed up the ordering.

[JackCaoG]

OK I think the issue is here

  partitions = partitioner.propose_partitions()
  print(partitions)

partitions doesn't seems to guarantee the order. Sometimes I see

[{add, index_copy_}]

other times I see

[{index_copy_, add}].   --> this one will produce the correct output

[JackCaoG]

If set allows_single_node_partition back to False, the test seems to consistently passing. @wonjoolee95 what's the reason you want to change that to true?

[wonjoolee95]

Thanks for the investigations, Jack. I saw that when allows_single_node_partition=False, the code path of this test case didn't actually fall through to the external_internal function. Only with this flag set to True, it reached the extract_internal function. Let me quickly double check that right now and get back.

[wonjoolee95]

Ok, I've set allows_single_node_partition back to False and the test/metrics look correct. And the code path does seem to properly enter into the extract_internal function. Now I can see the unit test consistently passing, and I can run and see that all the dynamo unit tests are passing:

/home/wonjoo/miniconda3/lib/python3.10/site-packages/torchvision/io/image.py:13: UserWarning: Failed to load image Python extension: 'libc10_cuda.so: cannot open shared object file: No such file or directory'If you don't plan on using image functionality from `torchvision.io`, you can ignore this warning. Otherwise, there might be something wrong with your environment. Did you have `libjpeg` or `libpng` installed before building `torchvision` from source?
  warn(
.
----------------------------------------------------------------------
Ran 11 tests in 27.713s

OK

This should be ready for review now.

[JackCaoG]

Let's hold on this pr until tmr, this pr also touches dynamo and I don't want it to break tmr's whl by accident, otherwise lGTM.

I am still concern about the partitoner might return ops with incorrect order, we should open a gihtub issue and follow up with Sherlock.