[JackCaoG]

This will help #4475

Currently each dynamo training step(without optimizer) will result in 5 graph execution

They are

1. sync_multi to sync input to fwd graph
2. Execute the fwd graph
3. sync_multi to sync input to bwd graph
4. Execute bwd graph
5. mark_step at the end of optimized_mod

My finding is that

1. The first graph executation is uecessary. It is triggered by GetXlaData in
   

   xla/torch_xla/csrc/init_python_bindings.cpp

   Lines 1700 to 1702 in 661ab6d

   	if (auto xla_tensor_ptr =
   	bridge::TryGetXlaTensor(ivalue.toTensor())) {
   	dataptr = xla_tensor_ptr->GetXlaData();

   
   However if I check the HLO, it is actually a DeviceData IR(not sure why). We don't need to call applyPendingGraph to trigger execeution to retrive the XLAData. Hence my fix in GetXlaData

2. The graph 3 and graph 5 are both triggered by aot_autograd. Instead of use xla_sync_multi to only sync the partial pending tensors and use mark_step to sync the rest, we can just do one mark_step in the beginning to combine 2 graphs. One thing to note is that we need to make first mark_step to be blocking.

With above 2 fixes, we can cut graph from 5 -> 3.

With the latest change, I am seeing a much more promising result
<style type="text/css"></style>

model	Old Speed Up	new Speed up
resnet50	0.758	0.937
resnet18	0.66	1.003
BERT_pytorch	1.441	1.869
resnext50_32x4d	0.87	1.139
alexnet	0.632	0.802
mobilenet_v2	0.549	0.672
mnasnet1_0	0.698	0.967
vgg16	0.712	0.742
timm_vision_transformer	1.275	1.69
squeezenet1_1	0.72	0.958
Avg	0.8315	1.0779

[JackCaoG]

I run the benchmark again using this pr and sadly find out the improvement is not as big as I would hope

The mark_step between forward and backward are very expensive, despite the actual device execution is very tiny. I will look into

1. If there is a way to make this mark_step faster
2. If there is a way to make this mark_step async

For the second part I actually had some idea. I might merge this pr as it is for now.

[JackCaoG]

<style type="text/css"></style>

model	Old Speed Up	New Speed Up
resnet50	0.758	0.731
resnet18	0.66	0.693
BERT_pytorch	1.441	1.459
resnext50_32x4d	0.87	0.845
alexnet	0.632	0.747
mobilenet_v2	0.549	0.546
mnasnet1_0	0.698	0.674
vgg16	0.712	0.731
timm_vision_transformer	1.275	1.335
squeezenet1_1	0.72	0.767
Avg	0.8315	0.8528	2.56% better

[shunting314]

In general, does a synchronous mark_step can be slower than _xla_sync_multi? Since the latter does not wait for the execution.

[JackCaoG]

_xla_sync_multi is also synchronous, but it might be executing a smaller graph.

[shunting314]

Ah, I didn't realize that the default behavior of 'mark_step (not wait for compute)' and '_xla_sync_multi (wait for compute)' are different. Thanks for clarification!

[JackCaoG]

Oh interesting, if I remove the mark_step at the end of the optimized_mod, resnet50 speed up jump from
0.731 -> 0.928

I think what happened is that even mark_step at the end might be a no-op, it will still try to acquire the device lock. However device lock is hold by the async thread of torch_xla._XLAC._run_cached_graph which blocks the mark_step.

I think we should not unconditionally perform the mark_step at the end

[JackCaoG]

Oh interesting, if I remove the mark_step at the end of the optimized_mod, resnet50 speed up jump from 0.731 -> 0.928

I think what happened is that even mark_step at the end might be a no-op, it will still try to acquire the device lock. However device lock is hold by the async thread of torch_xla._XLAC._run_cached_graph which blocks the mark_step.

I think we should not unconditionally perform the mark_step at the end

This also explains why my previous pr to make execution async for dynam doesn't help, it is because mark_step will block on that async execution anyway.

[alanwaketan]

LGTM.

[JackCaoG]

With the latest change, I am seeing a much more promising result
<style type="text/css"></style>

model	Old Speed Up	new Speed up
resnet50	0.758	0.937
resnet18	0.66	1.003
BERT_pytorch	1.441	1.869
resnext50_32x4d	0.87	1.139
alexnet	0.632	0.802
mobilenet_v2	0.549	0.672
mnasnet1_0	0.698	0.967
vgg16	0.712	0.742
timm_vision_transformer	1.275	1.69
squeezenet1_1	0.72	0.958
Avg	0.8315	1.0779

I think there still are some unit test failing, need to investigate a bit

[JackCaoG]

Ok I think I fixed all of the failing tests, let's wait for CI to confirm