Xla: Revisit mark_step() or torch_xla._XLAC._xla_sync_multi
torch_xla._XLAC._xla_sync_multi([c], devices=[], wait=True, sync_xla_data=True)
@dlibenzi Hi, I revisit the sync operation problem, wondering what's the real meaning behind the
sync op, this issue has been discussed in #2288 before (has been closed).
I modify your above code (in #2288) a little bit as following:
import time
import torch
import torch_xla
import torch_xla.core.xla_model as xm
m = 4096
n = 40928
k = 1024
h = min(n,k)
N = 201
print("torch version: ", torch.__version__)
device = xm.xla_device()
a = torch.randn(m,k).to(device)
b = torch.randn(k,n).to(device)
total_delta = 0
s = time.time()
for i in range(0, N):
c = torch.mm(a, b)
a[:, 0:h] = torch.min(a[:, 0:h], c[:, 0:h])
# xm.mark_step()
# print(torch_xla._XLAC._get_xla_tensors_text([c]))
xm.mark_step()
torch_xla._XLAC._xla_sync_multi([c], devices=[], wait=True, sync_xla_data=True)
total_delta = time.time() - s
c.cpu()
total_time = time.time() - s
print("TIME: ", total_time, total_delta)
FLOPS = 2.0 * (m * n *k) * (N - 1) / total_delta
print('{:.2f} GFLOPS'.format(FLOPS * 1e-9))
The output is: (the 391 TF/s seems too high)
torch version: 1.6.0
TIME: 4.922183036804199 0.17561125755310059
391010.18 GFLOPS
It seems using torch_xla._XLAC._xla_sync_multi or mark_step could not guarantee that
the loop has been executed.
From the IR report, we see exactly 201 aten::mm operations.
$grep "aten::mm" output.txt |wc
201 1006 10150
$ grep "aten::mm" output.txt |head
%10 = f32[4096,40928]{1,0} aten::mm(%9, %4)
%21 = f32[4096,40928]{1,0} aten::mm(%20, %4)
...
So we must depends on c.cpu() to execute the loop ?
Thanks.
_Originally posted by @shz0116 in https://github.com/pytorch/xla/issues/2288#issuecomment-699718881_
shz0116
All 21 comments
I run your code and checked that pending graph of c got materialized after the _xla_sync_multi so the graph is executed
IR {
%0 = f32[4096,40928]{1,0} xla::device_data(), location=<module>@test.py:28, device=TPU:0, ROOT=0
}
Looking at the compileTime, what's could happened was that TPU server side cache the graph so it does not recompile when you run your benchmarks.
Metric: XrtCompile
TotalSamples: 4
Accumulator: 05s331ms054.366us
Mean: 01s333ms763.591us
StdDev: 01s330ms961.608us
Rate: 0.0164502 / second
Percentiles: 25%=003ms083.051us; 50%=03s644ms395.886us; 80%=03s681ms929.050us; 90%=03s681ms929.050us; 95%=03s681ms929.050us; 99%=
03s681ms929.050us
Noted that there is a compile only takes 3ms while others takes 3s. You might want to restart your TPU between benchmarks to get a accurate result.
JackCaoG
on 28 Sep 2020
So I removed mark_step (leaving only _xla_sync_multi ) and saw
TIME: 6.914196729660034 3.85984468460083
after I restarted my TPU and
Metric: XrtCompile
TotalSamples: 1
Accumulator: 03s478ms269.453us
Mean: 03s478ms269.453us
StdDev: 000.000us
Percentiles: 25%=03s478ms269.453us; 50%=03s478ms269.453us; 80%=03s478ms269.453us; 90%=03s478ms269.453us; 95%=03s478ms269.453us; 9
9%=03s478ms269.453us
Metric: XrtExecute
TotalSamples: 1
Accumulator: 051ms106.581us
Mean: 051ms106.581us
StdDev: 000.000us
Percentiles: 25%=051ms106.581us; 50%=051ms106.581us; 80%=051ms106.581us; 90%=051ms106.581us; 95%=051ms106.581us; 99%=051ms106.581
us
which seems correct to me, delta_time is close to compile+execute
JackCaoG
on 28 Sep 2020
I noticed that HUGE time difference between the First run and the NEXT runs. Such as following.
And understand for the 2nd, 3rd runs, TPU cached the graph.
But I think we should used the results from 2rd or 3rd runs.
Since we want to measure the performance of torch.mm operation.
Building the execution graph is not our concern.
So, the results 390 TF/s or 380 TF/s is far higher than the theoretical peak 62 TF/s per core.
If we increase the iteration number N, we get even higher performance.
I am wondering how to explain the super good performance ?
I doubt some operations are optimized away even if I add data dependency.
Or other explanation ? @JackCaoG
Thanks,
(torch-xla-nightly) hongzhang@dlrm-tutorial:~/PytorchComputeKernels$ python3 he.py
torch version: 1.7.0a0+241afc9
TIME: 10.08446192741394 5.288620948791504
12983.69 GFLOPS
(torch-xla-nightly) hongzhang@dlrm-tutorial:~/PytorchComputeKernels$ python3 he.py
torch version: 1.7.0a0+241afc9
TIME: 4.521570920944214 0.17587757110595703
390418.11 GFLOPS
(torch-xla-nightly) hongzhang@dlrm-tutorial:~/PytorchComputeKernels$ python3 he.py
torch version: 1.7.0a0+241afc9
TIME: 4.584262847900391 0.1799478530883789
381587.16 GFLOPS
The only thing I can think of is that there might be some numerical simplification is applied during the compilation.
JackCaoG
on 28 Sep 2020
The only thing I can think of is that there might be some numerical simplification is applied during the compilation.
Is there a way to confirm ?
shz0116
on 28 Sep 2020
You can use the OP_BY_OP(doc) mode which will not perform the fusion of ops.
$ export XLA_GET_TENSORS_OPBYOP=1
$ export XLA_SYNC_TENSORS_OPBYOP=1
Execute time significantly increased while compile is much faster(no fusion)
Metric: CompileTime
TotalSamples: 1
Accumulator: 114ms403.740us
Percentiles: 1%=114ms403.740us; 5%=114ms403.740us; 10%=114ms403.740us; 20%=114ms403.740us; 50%=114ms403.740us; 80%=114ms403.740us; 90%=11
4ms403.740us; 95%=114ms403.740us; 99%=114ms403.740us
Metric: ExecuteChainedTime
TotalSamples: 1
Accumulator: 04s510ms489.727us
Percentiles: 1%=04s510ms489.727us; 5%=04s510ms489.727us; 10%=04s510ms489.727us; 20%=04s510ms489.727us; 50%=04s510ms489.727us; 80%=04s510m
s489.727us; 90%=04s510ms489.727us; 95%=04s510ms489.727us; 99%=04s510ms489.727us
and the TF is closer to what you would expect
TIME: 8.182087421417236 4.052510738372803
16944.01 GFLOPS
That being said, I am not sure if you want to benchmark this result, as OP_BY_OP is mostly used for the debugging purpose.
JackCaoG
on 28 Sep 2020
This is exactly I want to avoid. Otherwise I can move the sync into the loop body.
Trying to understand what's exactly happed during the execution process, leading to
the 390 TF/s.
shz0116
on 28 Sep 2020
Correct me if I'm wrong, but I see that you're only marking step / syncing at the end of the loop, which makes it only 1 execution; so all that the loop does is to produce the IR graph without execution, which is very very fast. And at the end, you only execute once. This is fine, but that would create one huge graph and cause problems (such as oom) as N grows.
For example, as I make
- N=201, I got ~520 TF/s
- N=501, I got ~826 TF/s
- N=1001, got ~705 TF/s
- N=4001, I got ~662 TF/s
- N=7001, I got ~559 TF/s
- N=10001, process errored.
Instead, when I move the sync inside the loop, I'm getting ~36TF/s.
Oh and I also added a warmup condition inside the loop, I don't measure the first 10 steps.
...
for i in range(0, N):
if i < warmup:
s = time.time()
...
So to answer your initial question; execution does not depend on c.cpu(). Removing that line does not affect TF/s for me at all.
taylanbil
on 29 Sep 2020
@taylanbil Hi, thanks for looking at this issue.
I can move the sync into the loop body as you mentioned. But when the matrix shape is small,
the results is very low. Such as m=128, n = 1024, k = 1024. I only got 1 TF/s.
That's why I am think moving sync out of loop body.
What I really confused is the FLOPS results reported, such as N=201, ~520 TF/s.
How the flops rate goes so high beyond the theoretical peak 62 TF/s per core (123 TF/s per chip)
on TPU-v3 ? That's why I start to think whether this is related with sync op.
Also saw an official announcement of Pytorch/XLA on TPU today. Congratulations.
shz0116
on 29 Sep 2020
Where does this theoretical peak come from?
With small inputs, xla will under-utilize the MXU's so it's normal that the performance suffers.
I am pretty sure that this is not related to syncing. As Jack pointed out, the hlo above c after marking step is a single device lookup, which means the execution is finished and data is on the device. To verify this, you could check the result of a computation for which the result is known. For example, instantiate matrices to be identity, and add one last sum like d = c.sum() before marking step, and then check the result and see if it's what you expect.
taylanbil
on 29 Sep 2020
Look at this paper Table 3 1st row PeakTeraflops 123 TF/s.
https://cacm.acm.org/magazines/2020/7/245702-a-domain-specific-supercomputer-for-training-deep-neural-networks/fulltext
shz0116
on 29 Sep 2020
yeah I'm confused by this as well. something is off somewhere.
taylanbil
on 29 Sep 2020
I tried to run the code in native cpu and found that after 30 steps, the value of c becomes nans
tensor([[-4.4271e+37, -1.3442e+38, 1.5414e+38, ..., 3.2558e+37,
-6.4465e+36, -1.4088e+38],
[-3.0807e+37, -8.5679e+37, 1.2530e+38, ..., 2.4030e+37,
-3.7335e+36, -1.4693e+38],
[-5.0853e+37, -9.1961e+37, 1.4596e+38, ..., 9.2441e+36,
-4.2325e+37, -1.4391e+38],
...,
[-3.3558e+37, -1.1097e+38, 1.1677e+38, ..., 2.7769e+37,
-1.9279e+37, -1.2275e+38],
[-3.4816e+36, -1.0689e+38, 1.4658e+38, ..., 4.3333e+37,
-9.4695e+36, -1.2794e+38],
[ 1.7933e+37, -8.8444e+37, 1.5581e+38, ..., 6.3365e+37,
4.5143e+36, -1.5187e+38]])
tensor([[-inf, nan, nan, ..., nan, nan, nan],
[-inf, -inf, nan, ..., nan, nan, nan],
[-inf, nan, nan, ..., nan, nan, nan],
...,
[-inf, nan, nan, ..., inf, nan, nan],
[-inf, nan, nan, ..., inf, nan, nan],
[-inf, nan, nan, ..., nan, nan, nan]])
tensor([[nan, nan, nan, ..., nan, nan, nan],
[nan, nan, nan, ..., nan, nan, nan],
[nan, nan, nan, ..., nan, nan, nan],
...,
[nan, nan, nan, ..., nan, nan, nan],
[nan, nan, nan, ..., nan, nan, nan],
[nan, nan, nan, ..., nan, nan, nan]])
I tested with these these inf or nan before, by default, tensorflow does not report any errors and just let it go.
I am not clear whether some optimization has been done for these special values.
A better initialization approach is use identity matrix as @taylanbil suggested to avoid these nans and inf.
shz0116
on 29 Sep 2020
I simplify your benchmark a bit
import time
import torch
import torch_xla
import torch_xla.core.xla_model as xm
m = 4096
N = 301
print("torch version: ", torch.__version__)
device = xm.xla_device()
a = torch.randn(m,m)
b = torch.randn(m,m)/ 10 + 1
a_xla = a.to(device)
b_xla = b.to(device)
total_delta = 0
s = time.time()
for i in range(0, N):
a_xla = torch.mm(a_xla, b_xla)
torch_xla._XLAC._xla_sync_multi([a_xla], devices=[], wait=True, sync_xla_data=True)
total_delta = time.time() - s
print(torch_xla._XLAC._get_xla_tensors_text([a_xla]))
import torch_xla.debug.metrics as met;
print(met.metrics_report())
a_xla.cpu()
total_time = time.time() - s
print("TIME: ", total_time, total_delta)
FLOPS = 2 * (m * m *m) * (N - 1) / total_delta
print('{:.2f} GFLOPS'.format(FLOPS * 1e-9))
The result still goes to intf pretty quick but I am consistently getting around 45~47TF with N=200 -> N=300
JackCaoG
on 30 Sep 2020
In your code:
a_xla = torch.mm(a_xla, b_xla)
a_xla will cause data dependence across iterations.
But this does not work for any matrix shape. such as a(m,k) * b (k,n) = c(m,.n), m,n, k, are different.
I need a "loop body" that work for any matrix shapes.
So I use:
a[:, 0:h] = torch.min(a[:, 0:h], c[:, 0:h])
to create data dependency across iterations to avoid compiler optimize the loop away.
So I could not use your change.
Also, could you explain why your loop body seems delivering "correct" results, but not my loop body ?
shz0116
on 30 Sep 2020
I am suspecting xla is smart enough to only calculate c[:, 0:h] instead of the whole matrix (since the rest of the matrix is not being used). This reduce the actual number of floating point calcuation.
JackCaoG
on 30 Sep 2020
This make a lot of sense.
So even pytorch/xla IR report 201 aten.mm, but each time, XLA compiler only compute "effective" matrix regions.
Could we contact google people to confirm this ?
shz0116
on 30 Sep 2020
I updated your code a bit and force xla to calculate the full matrix of c
import time
import torch
import torch_xla
import torch_xla.core.xla_model as xm
m = 4096
n = 40928
k = 1024
h = min(n,k)
N = 201
print("torch version: ", torch.__version__)
device = xm.xla_device()
a = torch.randn(m,k).to(device)
b = torch.randn(k,n).to(device)
c = torch.randn(m,n).to(device)
total_delta = 0
s = time.time()
for i in range(0, N):
c += torch.mm(a, b)
a[:, 0:h] = torch.min(a[:, 0:h], c[:, 0:h])
# xm.mark_step()
# print(torch_xla._XLAC._get_xla_tensors_text([c]))
xm.mark_step()
torch_xla._XLAC._xla_sync_multi([c], devices=[], wait=True, sync_xla_data=True)
total_delta = time.time() - s
c.cpu()
total_time = time.time() - s
print("TIME: ", total_time, total_delta)
FLOPS = ((m * n) * N + 2.0 * (m * n *k) * (N - 1) ) / total_delta
print('{:.2f} GFLOPS'.format(FLOPS * 1e-9))
The result looks reasonable now
(pytorch) jackcao@jack-dev:/tmp$ python test2.py
torch version: 1.7.0a0+c941dd3
TIME: 5.478667736053467 2.3137075901031494
29692.38 GFLOPS
Great. I'll give a test.
shz0116
on 30 Sep 2020
This issue has been automatically marked as stale because it has not had recent activity. It will be closed if no further activity occurs. Thank you for your contributions.
![stale[bot] picture](https://avatars3.githubusercontent.com/in/1724?v=4&s=40)
stale[bot]
on 1 Nov 2020
Related issues
ibeltagy
路
8Comments
david-alexander-white
路
6Comments
myleott
路
7Comments
magicknight
路
3Comments
ogulcanogul
路
5Comments
Most helpful comment
I am suspecting xla is smart enough to only calculate
c[:, 0:h]instead of the whole matrix (since the rest of the matrix is not being used). This reduce the actual number of floating point calcuation.