Debugging BERT Performance bottleneck

[anijain2305]

We are training BERT with PT-XLA. The script for training is present here - https://github.com/codeislife99/xla/blob/master/test/test_train_mp_bert_mlm.py

End to End results

	Iteration time (sec)	GPU time (sec)
PT-native	0.17	0.1654
PT-XLA	0.188	0.1105
Speedup PT-native vs PT-XLA	0.90426	1.49683

We observe these performance numbers. Iteration time is forward pass + backward pass + update.
In addition to iteration time, we use dlprof to measure the CUDA time. We observe that CUDA time is quite faster for PT-XLA (49% faster), but still the overall iteration time is 10% slower.

XRT Profile numbers

Name	Num_calls	Total time (ms)	Percentage	Acc Percentage
XrtCompile	6	266252.50842	78.55986	78.55986
XrtExecute	1280	72450.99722	21.37723	99.93709
XrtReleaseAllocation	17679	212.04111	0.06256	99.99965
XrtAllocateFromTensor	3416	0.69785	0.00021	99.99986
XrtReadLiteral	650	0.4893	0.00014	100

Auto-metric analysis

pt-xla-profiler: ================================================================================
pt-xla-profiler: Unlowered Op usage summary (more of these ops, lower performance)
pt-xla-profiler: Note: _local_scalar_dense typically indicates CPU context access
pt-xla-profiler: --------------------------------------------------------------------------------
pt-xla-profiler: FRAME (count=640):
pt-xla-profiler: Unlowered Op: "_local_scalar_dense"
pt-xla-profiler: Python Frames:
pt-xla-profiler:   <genexpr> (/root/anaconda3/envs/pytorch/lib/python3.6/site-packages/torch_xla/amp/grad_scaler.py:11)
pt-xla-profiler:   _maybe_opt_step (/root/anaconda3/envs/pytorch/lib/python3.6/site-packages/torch_xla/amp/grad_scaler.py:11)
pt-xla-profiler:   step (/root/anaconda3/envs/pytorch/lib/python3.6/site-packages/torch/cuda/amp/grad_scaler.py:339)
pt-xla-profiler:   loop_with_amp (/pytorch/xla/test/test_train_mp_bert_mlm.py:53)
pt-xla-profiler:   train (/pytorch/xla/test/test_train_mp_bert_mlm.py:165)
pt-xla-profiler:   main (/pytorch/xla/test/test_train_mp_bert_mlm.py:208)
pt-xla-profiler:   <module> (/pytorch/xla/test/test_train_mp_bert_mlm.py:237)
pt-xla-profiler:
pt-xla-profiler:
================================================================================

pt-xla-profiler: TransferFromServerTime too frequent: 648 counts during 1279 steps

• The CPU fallback to local_scalar_dense happens because of Step of GradScaler here - https://github.com/pytorch/xla/blob/master/torch_xla/amp/grad_scaler.py#L11. So, IIUC, this is unavoidable
• TransferFromServerTime is frequent because of local_scalar_dense, so it is a byproduct and not a new cause.

Manual Timing Analysis

We put timers in the scripts and code to measure different portions and found that forward and backward calls (which supposedly build the XLA graph) are taking substantial time

	Time (ms)
Forward	24
Backward	17

So, around 40-45 ms (out of 180 ms) in every iteration goes in the forward and backward calls.

@JackCaoG Does this analysis make sense? Can you provide any pointers?
@@codeislife99

[anijain2305]

Also, I added timestamp prints in the torch_xla/csrc/aten_xla_type.cpp for add and cmul operator. And I can see lots of prints in the forward pass call. This leads me to believe that it is actually building the XLA graph during that time. On subtracting the timestamps of first and last line, we can see that they are 26 ms apart. Is that expected?

2021-07-15 22:08:46.265262 Forward Starting
2021-07-15 22:08:46.265939: I   26933 torch_xla/csrc/aten_xla_type.cpp:463] An Op type add
2021-07-15 22:08:46.266336: I   26933 torch_xla/csrc/aten_xla_type.cpp:525] An Op type cmul
2021-07-15 22:08:46.267885: I   26933 torch_xla/csrc/aten_xla_type.cpp:463] An Op type add
2021-07-15 22:08:46.268417: I   26933 torch_xla/csrc/aten_xla_type.cpp:463] An Op type add
2021-07-15 22:08:46.268544: I   26933 torch_xla/csrc/aten_xla_type.cpp:525] An Op type cmul
2021-07-15 22:08:46.269102: I   26933 torch_xla/csrc/aten_xla_type.cpp:463] An Op type add
2021-07-15 22:08:46.269226: I   26933 torch_xla/csrc/aten_xla_type.cpp:525] An Op type cmul
2021-07-15 22:08:46.269853: I   26933 torch_xla/csrc/aten_xla_type.cpp:463] An Op type add
2021-07-15 22:08:46.270369: I   26933 torch_xla/csrc/aten_xla_type.cpp:463] An Op type add
2021-07-15 22:08:46.270494: I   26933 torch_xla/csrc/aten_xla_type.cpp:525] An Op type cmul
2021-07-15 22:08:46.270996: I   26933 torch_xla/csrc/aten_xla_type.cpp:463] An Op type add
2021-07-15 22:08:46.271122: I   26933 torch_xla/csrc/aten_xla_type.cpp:525] An Op type cmul
2021-07-15 22:08:46.271757: I   26933 torch_xla/csrc/aten_xla_type.cpp:463] An Op type add
2021-07-15 22:08:46.272306: I   26933 torch_xla/csrc/aten_xla_type.cpp:463] An Op type add
2021-07-15 22:08:46.272434: I   26933 torch_xla/csrc/aten_xla_type.cpp:525] An Op type cmul
2021-07-15 22:08:46.272990: I   26933 torch_xla/csrc/aten_xla_type.cpp:463] An Op type add
2021-07-15 22:08:46.273115: I   26933 torch_xla/csrc/aten_xla_type.cpp:525] An Op type cmul
2021-07-15 22:08:46.273762: I   26933 torch_xla/csrc/aten_xla_type.cpp:463] An Op type add
2021-07-15 22:08:46.274294: I   26933 torch_xla/csrc/aten_xla_type.cpp:463] An Op type add
2021-07-15 22:08:46.274422: I   26933 torch_xla/csrc/aten_xla_type.cpp:525] An Op type cmul
2021-07-15 22:08:46.274937: I   26933 torch_xla/csrc/aten_xla_type.cpp:463] An Op type add
2021-07-15 22:08:46.275064: I   26933 torch_xla/csrc/aten_xla_type.cpp:525] An Op type cmul
2021-07-15 22:08:46.275715: I   26933 torch_xla/csrc/aten_xla_type.cpp:463] An Op type add
2021-07-15 22:08:46.276256: I   26933 torch_xla/csrc/aten_xla_type.cpp:463] An Op type add
2021-07-15 22:08:46.276383: I   26933 torch_xla/csrc/aten_xla_type.cpp:525] An Op type cmul
2021-07-15 22:08:46.276946: I   26933 torch_xla/csrc/aten_xla_type.cpp:463] An Op type add
2021-07-15 22:08:46.277077: I   26933 torch_xla/csrc/aten_xla_type.cpp:525] An Op type cmul
2021-07-15 22:08:46.277729: I   26933 torch_xla/csrc/aten_xla_type.cpp:463] An Op type add
2021-07-15 22:08:46.278274: I   26933 torch_xla/csrc/aten_xla_type.cpp:463] An Op type add
2021-07-15 22:08:46.278401: I   26933 torch_xla/csrc/aten_xla_type.cpp:525] An Op type cmul
2021-07-15 22:08:46.278914: I   26933 torch_xla/csrc/aten_xla_type.cpp:463] An Op type add
2021-07-15 22:08:46.279044: I   26933 torch_xla/csrc/aten_xla_type.cpp:525] An Op type cmul
2021-07-15 22:08:46.279693: I   26933 torch_xla/csrc/aten_xla_type.cpp:463] An Op type add
2021-07-15 22:08:46.280238: I   26933 torch_xla/csrc/aten_xla_type.cpp:463] An Op type add
2021-07-15 22:08:46.280367: I   26933 torch_xla/csrc/aten_xla_type.cpp:525] An Op type cmul
2021-07-15 22:08:46.280921: I   26933 torch_xla/csrc/aten_xla_type.cpp:463] An Op type add
2021-07-15 22:08:46.281049: I   26933 torch_xla/csrc/aten_xla_type.cpp:525] An Op type cmul
2021-07-15 22:08:46.281695: I   26933 torch_xla/csrc/aten_xla_type.cpp:463] An Op type add
2021-07-15 22:08:46.282248: I   26933 torch_xla/csrc/aten_xla_type.cpp:463] An Op type add
2021-07-15 22:08:46.282376: I   26933 torch_xla/csrc/aten_xla_type.cpp:525] An Op type cmul
2021-07-15 22:08:46.282886: I   26933 torch_xla/csrc/aten_xla_type.cpp:463] An Op type add
2021-07-15 22:08:46.283015: I   26933 torch_xla/csrc/aten_xla_type.cpp:525] An Op type cmul
2021-07-15 22:08:46.283662: I   26933 torch_xla/csrc/aten_xla_type.cpp:463] An Op type add
2021-07-15 22:08:46.284204: I   26933 torch_xla/csrc/aten_xla_type.cpp:463] An Op type add
2021-07-15 22:08:46.284333: I   26933 torch_xla/csrc/aten_xla_type.cpp:525] An Op type cmul
2021-07-15 22:08:46.284897: I   26933 torch_xla/csrc/aten_xla_type.cpp:463] An Op type add
2021-07-15 22:08:46.285025: I   26933 torch_xla/csrc/aten_xla_type.cpp:525] An Op type cmul
2021-07-15 22:08:46.285663: I   26933 torch_xla/csrc/aten_xla_type.cpp:463] An Op type add
2021-07-15 22:08:46.286196: I   26933 torch_xla/csrc/aten_xla_type.cpp:463] An Op type add
2021-07-15 22:08:46.286321: I   26933 torch_xla/csrc/aten_xla_type.cpp:525] An Op type cmul
2021-07-15 22:08:46.286829: I   26933 torch_xla/csrc/aten_xla_type.cpp:463] An Op type add
2021-07-15 22:08:46.286957: I   26933 torch_xla/csrc/aten_xla_type.cpp:525] An Op type cmul
2021-07-15 22:08:46.287586: I   26933 torch_xla/csrc/aten_xla_type.cpp:463] An Op type add
2021-07-15 22:08:46.288105: I   26933 torch_xla/csrc/aten_xla_type.cpp:463] An Op type add
2021-07-15 22:08:46.288231: I   26933 torch_xla/csrc/aten_xla_type.cpp:525] An Op type cmul
2021-07-15 22:08:46.288773: I   26933 torch_xla/csrc/aten_xla_type.cpp:463] An Op type add
2021-07-15 22:08:46.288898: I   26933 torch_xla/csrc/aten_xla_type.cpp:525] An Op type cmul
2021-07-15 22:08:46.289527: I   26933 torch_xla/csrc/aten_xla_type.cpp:463] An Op type add
2021-07-15 22:08:46.290044: I   26933 torch_xla/csrc/aten_xla_type.cpp:463] An Op type add
2021-07-15 22:08:46.290168: I   26933 torch_xla/csrc/aten_xla_type.cpp:525] An Op type cmul
2021-07-15 22:08:46.290665: I   26933 torch_xla/csrc/aten_xla_type.cpp:463] An Op type add
2021-07-15 22:08:46.290788: I   26933 torch_xla/csrc/aten_xla_type.cpp:525] An Op type cmul
2021-07-15 22:08:46.291203: I   26933 torch_xla/csrc/aten_xla_type.cpp:525] An Op type cmul
2021-07-15 22:08:46.291597 Forward finished

[JackCaoG]

Sorry for the late reply, the metric actually seems pretty reasonable to me. Compilation take a long time(compared to the Execution) but it should stabilize after first few epoch(no more compilation and only execution). When you say the overall iteration time is 10% slower how many epoch did you test? You can image that if after first epochs, no compilation happens, the step time will be much faster for the following epochs.

Also regarding the grad_scalar trigger a roundtrip to cpu, let me spend sometime to see if we can avoid this.

[anijain2305]

@JackCaoG Thanks for the reply. The above numbers are reported after second epoch. I have already removed the first epoch, where compilation happens. Therefore, we observe 10% slowdown after second epoch (first epoch is much worse, it takes > 1m to finish compilation).

The above time I am pointing towards is potentially just parsing the PyTorch ops to XLA ops. The compilation of the resulting XLA graph is already done in the first epoch, so no compilation here.

[anijain2305]

I also varied the batch size to get some more confidence that it is indeed parsing time. Again, this is after second epoch. I observe that graph creation time is more or less same for different batch sizes (which is expected).

Batch size = 8
Category	Time (ms)	Percent (%)	Accumated Perc (%)
Forward graph creation time	24	12.9	12.9
Backward graph creation time	17	9.14	22.04
CUDA time	118	63.44	85.48
Unknown	27	14.52	100
Total time	186
Batch size = 12
Category	Time (ms)	Percent (%)	Accumated Perc (%)
Forward graph creation time	23	10.31	10.31
Backward graph creation time	21	9.42	19.73
CUDA time	147	65.92	85.65
Unknown	32	14.35	100
Total time	223
Batch size = 16
Category	Time (ms)	Percent (%)	Accumated Perc (%)
Forward graph creation time	23	8.81	8.81
Backward graph creation time	21	8.05	16.86
CUDA time	194	74.33	91.19
Unknown	23	8.81	100
Total time	261

---

Finally, the performance comparison looks like this. As batch size increases, the parsing time remains constant and its contribution to overall iteration time reduces. This helps improve e2e speedup as XLA improves CUDA time.

	Iteration Time (sec)
Batch size	PT-native	PT-XLA	Speedup of PT-XLA
4	0.104	0.141	0.73759
8	0.17	0.186	0.91398
12	0.234	0.223	1.04933
16	FAILS	0.261	#VALUE!

[JackCaoG]

Thanks for the detail analysis! I think you are right that bottleneck here is the (IR) graph building time, we will do that every epoch but we won't perform the IR -> HLO -> executable compilation unless IR graph changed. This can become a burden when execution time is relatively short. I don't think there is much we can do about the graph building since pt/xla is an eager runtime after all. I will try to look at grad_scalar issue but it seems like unless we also changed the pytorch frontend API (currently takes None if inf is founded which will force us to eval the graph early) we have to stick with the current approach.

It seems like PT/XLA scales better than native pytorch which is a good sign.

[anijain2305]

@JackCaoG Thanks for the reply. Yes, my observation is similar. XLA compilation happens only in early iterations. But, IR graph creation happens in every iteration.

Just want to confirm one more thing on your statement

I think you are right that bottleneck here is the (IR) graph building time, we will do that every epoch but we won't perform the IR -> HLO -> executable compilation unless IR graph changed

When you say epoch here, do you mean iteration? My analysis shows that (IR) graph building happens in every iteration OR specifically every time we make forward() and backward() calls. Just wanted to confirm if you also meant the same.

[qinggangz]

Thanks for the suggestions, @JackCaoG. "we won't perform the IR -> HLO -> executable compilation unless IR graph changed" Is there an easy way to determine if IR graph changed without creating it? It seems 8%~12% overhead is for creating IR graph alone.

[byronyi]

I think you are right that bottleneck here is the (IR) graph building time, we will do that every epoch but we won't perform the IR -> HLO -> executable compilation unless IR graph changed

I believe @JackCaoG means that we build the IR graph every training step. I doubt we have any API to inform the tracer that the IR graph won't change across each step. Maybe we could add one, like pytorch/pytorch#15623?

[byronyi]

It would also be helpful to add an utility to detect whether the IR graph changes, much like check_inputs for torch.jit.trace:

def loop_in_traced_fn(x):
    result = x[0]
    for i in range(x.size(0)):
        result = result * x[i]
    return result

inputs = (torch.rand(3, 4, 5),)
check_inputs = [(torch.rand(4, 5, 6),), (torch.rand(2, 3, 4),)]

traced = torch.jit.trace(loop_in_traced_fn, inputs, check_inputs=check_inputs)

When this occurs, we could fallback to pure eager runtime while prints a nice error like the following for TS:

ERROR: Graphs differed across invocations!
Graph diff:

            graph(%x : Tensor) {
            %1 : int = prim::Constant[value=0]()
            %2 : int = prim::Constant[value=0]()
            %result.1 : Tensor = aten::select(%x, %1, %2)
            %4 : int = prim::Constant[value=0]()
            %5 : int = prim::Constant[value=0]()
            %6 : Tensor = aten::select(%x, %4, %5)
            %result.2 : Tensor = aten::mul(%result.1, %6)
            %8 : int = prim::Constant[value=0]()
            %9 : int = prim::Constant[value=1]()
            %10 : Tensor = aten::select(%x, %8, %9)
        -   %result : Tensor = aten::mul(%result.2, %10)
        +   %result.3 : Tensor = aten::mul(%result.2, %10)
        ?          ++
            %12 : int = prim::Constant[value=0]()
            %13 : int = prim::Constant[value=2]()
            %14 : Tensor = aten::select(%x, %12, %13)
        +   %result : Tensor = aten::mul(%result.3, %14)
        +   %16 : int = prim::Constant[value=0]()
        +   %17 : int = prim::Constant[value=3]()
        +   %18 : Tensor = aten::select(%x, %16, %17)
        -   %15 : Tensor = aten::mul(%result, %14)
        ?     ^                                 ^
        +   %19 : Tensor = aten::mul(%result, %18)
        ?     ^                                 ^
        -   return (%15);
        ?             ^
        +   return (%19);
        ?             ^
            }

Ultimately we would like a better prompt message, so the users could know if the dynamism is an essential part of the model or not, and if it isn't, how to change their Python code to make it a static graph.

[JackCaoG]

yup, I meant every step. It is also common to see IR being built every forward and backward call.

Currently there is no way to skip the IR graph building process. We compute the hash of the IR graph as we are building it and we use hash to determine if two graph is the same.

Tracing a graph once and run it many time is actually not easy for the eager frontend. For the most of XlaTensor we have a IRValue(most common case, can also be a view or device data, depending on the tensor) associated with it. When the mark_step is called, we will compile(if cache miss) and execute the graph, then replacing the IrValue of the XlaTensor with the actual device data pointer. The process of building the IR graph link each tensor with the result of the computation. Without too much thinking, removing the IR building process will make accessing intermediate XlaTensor pretty difficult. I am sure there are more problems with this approach.

[byronyi]

Or we could do a little bit of branch prediction, i.e. let the frontend tracing and runtime execution occur concurrently, assuming the IR graph is not changing. When it changes though, we cancel the in-progress runtime execution, print a warning that in which PyTorch kernel dispatch the IR graph starts to diverge which corresponds directly to the Python code in user land.

This would be more transparent to users as it doesn't need additional API calls or explicit annotation (things could go south as most users simply assume their graph aren't changing appearantly). The tracer in the frontend has better context for a potentially nice warning message.

[JackCaoG]

@byronyi We do the feature of detecting IR graph changed, if you do a debug run and there will be a graph report. Within the report, you will see something like

  Frame 5 (len=2606, count=1, id=55, h=(5620df66425d34129b9d87de628db35a)) vs 6 (len=2032, count=1, id=56 h=(979909\
444b27e0a47d62d5643d6a5aad))                                                                                        
  --- frame-5                                                                                                       
  +++ frame-6                                                                                                       
  @@ -1,10 +1,24 @@                                                                                                 
   IR {                                                                                                             
  -  f32[] xla::device_data(), location=calculate_adaptive_weight@vqperceptual.py:83, device=TPU:0                  
  -  f32[] xla::device_data(), location=calculate_adaptive_weight@vqperceptual.py:82, device=TPU:0                  
  -  f32[] prim::Constant(), location=calculate_adaptive_weight@vqperceptual.py:82, value=0                         
  -  f32[] prim::Constant(), location=calculate_adaptive_weight@vqperceptual.py:81, value=1                         
  -  f32[] xla::device_data(), location=calculate_adaptive_weight@vqperceptual.py:81, device=TPU:0                  
  -  f32[] aten::mul(?, ?), location=calculate_adaptive_weight@vqperceptual.py:81                                   
  +  f32[1]{0} xla::device_data(), location=_conv_forward@conv.py:439, device=TPU:0                                 
  +  f32[1,512,4,4]{1,3,2,0} xla::device_data(), location=_conv_forward@conv.py:439, device=TPU:0                   
  +  f32[512]{0} xla::device_data(), location=batch_norm@functional.py:2281, device=TPU:0                           
  +  f32[512]{0} xla::device_data(), location=batch_norm@functional.py:2281, device=TPU:0 
.....

[anijain2305]

Is it possible to overlap the IR graph creation of step i with the kernel execution of step i -1?