[Functionalization] dynamo test failure

[wonjoo-wj]

With the latest rebases of the functionalization branches, we're seeing new test failure related to dynamo:

(base) jenkins@26d7adccbc26:/workspace/pytorch/xla$ python test/dynamo/test_dynamo.py -k test_resnet18
F
======================================================================
FAIL: test_resnet18 (__main__.DynamoTrainingBasicTest)
----------------------------------------------------------------------
Traceback (most recent call last):
  File "test/dynamo/test_dynamo.py", line 163, in test_resnet18
    self.assertEqual(met.metric_data('CompileTime')[0], 4)
AssertionError: 5 != 4

----------------------------------------------------------------------
Ran 2 tests in 41.536s

FAILED (failures=1)

Previously, we had dynamo tests fail that were failing due to seemingly related metric assertions (issue -- #4414). This error is different, as it does not throw any additional warnings and only the graph compilation metric counter comparison is failing.

Creating this issue for tracking purposes.

[JackCaoG]

you can dump the graph before and after functionizaiton then it should be obvious where the additional graph is from.

[wonjoo-wj]

Putting some IR dumps here. With master branch, dumping the graph of xla_output at https://github.com/pytorch/xla/blob/master/test/dynamo/test_dynamo.py#L146:

<master branch>

(base) jenkins@26d7adccbc26:/workspace/pytorch/xla$ python test/dynamo/test_dynamo.py -k test_resnet18
IR {
  %0 = f32[4,1000]{1,0} xla::device_data(), location=test_resnet18@test_dynamo.py:148, xla_shape=f32[4,1000]{1,0}, device=CPU:0, ROOT=0
}
...
IR {
  %0 = f32[4,1000]{1,0} xla::device_data(), location=test_resnet18@test_dynamo.py:148, xla_shape=f32[4,1000]{1,0}, device=CPU:0, ROOT=0
}
Ran 2 tests in 88.218s

OK
(base) jenkins@26d7adccbc26:/workspace/pytorch/xla$ 

Note that the xla_output is in a loop for data, target in loader hence the multiple IR graphs. The dumps in the functionalization branch actually are equal:

<functionalization branch>

(base) jenkins@26d7adccbc26:/workspace/pytorch/xla$ python test/dynamo/test_dynamo.py -k test_resnet18
IR {
  %0 = f32[4,1000]{1,0} xla::device_data(), location=test_resnet18@test_dynamo.py:148, xla_shape=f32[4,1000]{1,0}, device=CPU:0, ROOT=0
}
...
IR {
  %0 = f32[4,1000]{1,0} xla::device_data(), location=test_resnet18@test_dynamo.py:148, xla_shape=f32[4,1000]{1,0}, device=CPU:0, ROOT=0
}

F
======================================================================
FAIL: test_resnet18 (__main__.DynamoTrainingBasicTest)
----------------------------------------------------------------------
Traceback (most recent call last):
  File "test/dynamo/test_dynamo.py", line 161, in test_resnet18
    self.assertEqual(met.metric_data('CompileTime')[0], 4)
AssertionError: 5 != 4

----------------------------------------------------------------------
Ran 2 tests in 37.397s

FAILED (failures=1)
(base) jenkins@26d7adccbc26:/workspace/pytorch/xla$ 

[JackCaoG]

It is more than just output because it is a training so we will have forward + backward + some random clean up graph. You want to use things like XLA_SAVE_TENSORS_FMT='text' XLA_SAVE_TENSORS_FILE='/tmp/save1.ir' to dump all graph being generated. Note that graph dump with [ScheduleSyncTensorsGraph] means it is caused by a mark_step, graph dump with [dynamo] means it is a graph cached and we execute it with dynamo.

I suspect you will see one additional graph with functionization branch

[wonjoo-wj]

Thanks for the pointers, Jack. Posting some more dumps here. Ran the master and functionalization branch:

XLA_SAVE_TENSORS_FMT='text' XLA_SAVE_TENSORS_FILE='/workspace/functionalization-dynamo.ir' python test/dynamo/test_dynamo.py -k test_resnet18

The IR dump is quite large as this a training:

• master: https://gist.github.com/wonjoolee95/c7ecd644dab680c1e8d72a74d4ccc4d4
• functionalization: https://gist.github.com/wonjoolee95/d277f8d401de9bce38c7c93c3100fe11

Couple of things to note:

• Just doing a quick search on dynamo graph [dynamo], the master branch dump shows 29 matches whereas the functionalization branch dump shows 40 matches.
• For [ScheduleSyncTensorsGraph], the master branch dump shows 15 matches whereas the functionalization branch dump shows 17 matches.
• The training loop (at https://github.com/pytorch/xla/blob/master/test/dynamo/test_dynamo.py#L145) defaults to 10 iterations as SAMPLE_COUNT=10, hence the huge dump files/graphs.

[JackCaoG]

hmm @wonjoolee95 let's comment out the CompileTime check and see how many execution it actually performed. reduce the iteration from 10 to 3 maybe, to make debugging easier.

[wonjoo-wj]

Reducing the iteration to 3, and the metrics for the functionalization branch looks like:

METRICS: met.metric_data('CompileTime')[0]: 5
METRICS: met.metric_data('ExecuteTime')[0]: 9
METRICS: met.metric_data('RunCachedGraphInputData')[0]: 6
METRICS: met.metric_data('RunCachedGraphOutputData')[0]: 6

Whereas the master branch looks like:

METRICS: met.metric_data('CompileTime')[0]: 4
METRICS: met.metric_data('ExecuteTime')[0]: 9
METRICS: met.metric_data('RunCachedGraphInputData')[0]: 6
METRICS: met.metric_data('RunCachedGraphOutputData')[0]: 6

Other than CompileTime, the other metrics seem to be the same/correct. The full IR dump for the 3-iteration versions:

• master: https://gist.github.com/wonjoolee95/ff3c01e9637c844301f3b3599ac704d0
• functionalization: https://gist.github.com/wonjoolee95/9ba4422ebbfc32b58bb0eac91a74e555

[JackCaoG]

lol OK this is a bit hard to read, pretty you need to identify all of the unique graphs with [ScheduleSyncTensorsGraph] and find out the additional compile..

I am considered to add a unique_compile flag to make this kind of debugging eaiser...

[wonjoo-wj]

FWIW -- just going through those two dumps, the first deviation starts at the 6th [ScheduleSyncTensorsGraph] at https://gist.github.com/wonjoolee95/ff3c01e9637c844301f3b3599ac704d0#file-master-dynamo-3-L3767 (for master) and https://gist.github.com/wonjoolee95/9ba4422ebbfc32b58bb0eac91a74e555#file-functionalization-dynamo-3-L3767 (for functionalization).

The IR graph starting from this line deviates -- the 0% entry for the functionalization branch for some reason calls meta_convolution_backward:

  %0 = f32[] prim::Constant(), location=meta_convolution_backward@_meta_registrations.py:1121, xla_shape=f32[], value=0
  %1 = f32[64,3,7,7]{3,2,1,0} aten::expand(%0), location=meta_convolution_backward@_meta_registrations.py:1121, xla_shape=f32[64,3,7,7]{3,2,1,0}, size=(64, 3, 7, 7), ROOT=0

As opposed to the master branch:

  %0 = f32[64,3,7,7]{3,2,1,0} xla::device_data(), location=__call__@_ops.py:284, xla_shape=f32[64,3,7,7]{3,2,1,0}, device=CPU:0

[alanwaketan]

@wonjoolee95 Hmm, I'm working on fixing one of our conv issue: #4681. Wondering if that could solve your problem here.

[alanwaketan]

The DynamoTrainingOptimizerTest.test_resnet18 is broken too.

[JackCaoG]

That's OK, it is likely the same failure since it is a superset of training test.

[wonjoo-wj]

Okay coming back to this, still a mess to go through these dumps lol.

One new thing I did notice, though, is that when the iteration (SAMPLE_COUNT) is reduced to 2, the test metrics are correct:

WONJOO: met.metric_data('CompileTime')[0]: 4
WONJOO: met.metric_data('ExecuteTime')[0]: 6
WONJOO: met.metric_data('RunCachedGraphInputData')[0]: 4
WONJOO: met.metric_data('RunCachedGraphOutputData')[0]: 4

The CompileTime metric starts to be 5 when SAMPLE_COUNT >= 3.