Xla: Revisit how RNG seed moves forward.

[taylanbil]

🚀 Feature

RNG seed moving forward on its own, without requiring data transfer from the client VM, would speed workloads that use dropouts and random operations much faster.

Motivation

Currently, the seed on the device is reset every time we mark step here, and then we get the new rng seed for the new step from the client, here.

This is causing a large idle time for workloads that use dropout. Example: Roberta, transformer, etc. For some cases (of high interest), this causes TransferToServerTime to be about as high as ExecuteTime, and consequently TPU cores sit idle > 40% of the time.

Pitch

The seed should move forward on the device, without requiring an input per step from the client, thus not causing high TransferToServerTime.

Additional context

It would be also useful to fork the RNG state, and revert back to where it was, to do dropout correctly in Model-Parallel workloads. This is possible to do, thanks to https://github.com/pytorch/xla/pull/2096, but this too requires data to be transferred from client to server. It would be great to not need that transfer to achieve this. I'm adding this context because I believe it may help us choose between possible approaches.

[taylanbil]

cc: @davidel @JackCaoG @myleott

[taylanbil]

Some thoughts around this, from discussions w/ Jack:

• Seed on device is not a tensor, it's a scalar. We're not sure if modifying the scalars by a mul+add is possible.
• Making the seed a tensor would allow it to be modifiable, but it would break the 1-1 correspondence b/w tensors on the client side vs server side.

[JackCaoG]

@davidel I am wondering if instead of saving seed_ir_value in dev_ctx, we save a seed_xla_tensor in dev_ctx and keep updating its ir_value to the current seed. This way when the step is mark, the value of this seed_xla_tensor will be updated. We don't need to manually update base_seed every step.

I am thinking something along the line of

@@ -268,6 +271,8 @@ class XLATensor::DeviceContextArena {
     xla::uint64 seed = 101;
     xla::uint64 running_seed = 101;
     ir::Value seed_ir_value;
+    XLATensor seed_tensor;
+    bool tensor_initialized = false;
   };

  public:
@@ -311,9 +316,10 @@ class XLATensor::DeviceContextArena {
     static const xla::uint64 kSeedAdd = 2531011;
     DeviceContext* devctx = GetDeviceContext(device);
     std::lock_guard<std::mutex> lock(devctx->lock);
-    if (!devctx->seed_ir_value) {
-      devctx->seed_ir_value =
-          IrValueFromScalar(MakeIntScalar(devctx->seed), kSeedType, device);
+    if (!devctx->tensor_initialized || !devctx->seed_tensor.CurrentIrValue()) {
+      ir::Value val = IrValueFromScalar(MakeIntScalar(devctx->seed), kSeedType, device);
+      devctx->seed_tensor = XLATensor::Create(val, device);
+      devctx->tensor_initialized = true;
     }
     // Keep the running seed as scalar as well, so we can return it directly
     // without executing graphs.
@@ -324,8 +330,8 @@ class XLATensor::DeviceContextArena {
                                     MakeXlaPrimitiveType(kSeedType, &device));
     ir::Value b = ir::ops::ScalarOp(MakeIntScalar(kSeedAdd),
                                     MakeXlaPrimitiveType(kSeedType, &device));
-    devctx->seed_ir_value = b + k * devctx->seed_ir_value;
-    return devctx->seed_ir_value;
+    devctx->seed_tensor.SetIrValue(b + k * devctx->seed_ir_value);
+    return devctx->seed_tensor.CurrentIrValue();
   }

   xla::uint64 GetRunningSeed(const Device& device) {
@@ -339,15 +345,11 @@ class XLATensor::DeviceContextArena {
     std::lock_guard<std::mutex> lock(devctx->lock);
     devctx->seed = seed;
     devctx->running_seed = devctx->seed;
-    devctx->seed_ir_value = ir::Value();
+    devctx->seed_tensor.SetIrValue(ir::Value());;
   }

   void MarkStep(const Device& device) {
-    DeviceContext* devctx = GetDeviceContext(device);
-    std::lock_guard<std::mutex> lock(devctx->lock);
-    devctx->seed = 1012031 + devctx->seed * 7012063;
-    devctx->running_seed = devctx->seed;
-    devctx->seed_ir_value = ir::Value();
+    return;
   }

The above change will hang the program so I think either I made a mistake somewhere or this approach fundamentally doesn't work 😄 ..

[davidel]

There seem to be something odd going on.
There should be at most one scalar transfer per step, which should be in the order of 1..2ms.
Every other seed handed over within a single step (ie, many dropouts), is composed by in-graph add+mul, which do not cause any transfers.
Unless, of course, the model being debugged has a step time of 1..2ms, in which case such time becomes non negligible.

My guess is, the code under scrutiny manually reset the seed before creating new dropouts layers, which is bogus.
But a simple test case with a few dropout layers will reveal the eventual issue, and help debugging.

[davidel]

I did a simple test:

import torch
import torch.nn as nn
import torch_xla
import torch_xla.core.xla_model as xm

device = xm.xla_device()

d1 = nn.Dropout(0.2)
d2 = nn.Dropout(0.25)
d3 = nn.Dropout(0.33)

t = torch.randn(3, 3, device=device)

t1 = d1(t)
t2 = d2(t1)
t3 = d3(t2)

print(torch_xla._XLAC._get_xla_tensors_text([t3]))

print(t3.cpu())

Which result in this:

IR {
  %0 = f32[] xla::device_data(), [email protected]:958, device=CPU:0
  %1 = s64[] xla::device_data(), location=<module>@dropout_test.py:12, device=CPU:0
  %2 = s64[] prim::Constant(), location=<module>@dropout_test.py:12, value=214013
  %3 = s64[] aten::mul(%2, %1), location=<module>@dropout_test.py:12
  %4 = s64[] prim::Constant(), location=<module>@dropout_test.py:12, value=2.53101e+06
  %5 = s64[] aten::add(%4, %3), location=<module>@dropout_test.py:12
  %6 = s64[] prim::Constant(), [email protected]:958, value=214013
  %7 = s64[] aten::mul(%6, %5), [email protected]:958
  %8 = s64[] prim::Constant(), [email protected]:958, value=2.53101e+06
  %9 = s64[] aten::add(%8, %7), [email protected]:958
  %10 = s64[] prim::Constant(), [email protected]:958, value=214013
  %11 = s64[] aten::mul(%10, %9), [email protected]:958
  %12 = s64[] prim::Constant(), [email protected]:958, value=2.53101e+06
  %13 = s64[] aten::add(%12, %11), [email protected]:958
  %14 = s64[] prim::Constant(), [email protected]:958, value=214013
  %15 = s64[] aten::mul(%14, %13), [email protected]:958
  %16 = s64[] prim::Constant(), [email protected]:958, value=2.53101e+06
  %17 = s64[] aten::add(%16, %15), [email protected]:958
  %18 = f32[] xla::device_data(), [email protected]:958, device=CPU:0
  %19 = f32[3,3]{1,0} aten::expand(%18), [email protected]:958, size=(3, 3)
  %20 = f32[3,3]{1,0} aten::bernoulli(%19, %17), [email protected]:958
  %21 = f32[3,3]{1,0} aten::div(%20, %0), [email protected]:958
  %22 = f32[] xla::device_data(), [email protected]:958, device=CPU:0
  %23 = f32[] xla::device_data(), [email protected]:958, device=CPU:0
  %24 = f32[3,3]{1,0} aten::expand(%23), [email protected]:958, size=(3, 3)
  %25 = f32[3,3]{1,0} aten::bernoulli(%24, %13), [email protected]:958
  %26 = f32[3,3]{1,0} aten::div(%25, %22), [email protected]:958
  %27 = f32[] xla::device_data(), [email protected]:958, device=CPU:0
  %28 = f32[] xla::device_data(), [email protected]:958, device=CPU:0
  %29 = f32[3,3]{1,0} aten::expand(%28), [email protected]:958, size=(3, 3)
  %30 = f32[3,3]{1,0} aten::bernoulli(%29, %9), [email protected]:958
  %31 = f32[3,3]{1,0} aten::div(%30, %27), [email protected]:958
  %32 = f32[] prim::Constant(), location=<module>@dropout_test.py:12, value=1
  %33 = f32[3,3]{1,0} aten::expand(%32), location=<module>@dropout_test.py:12, size=(3, 3)
  %34 = f32[] prim::Constant(), location=<module>@dropout_test.py:12, value=0
  %35 = f32[3,3]{1,0} aten::expand(%34), location=<module>@dropout_test.py:12, size=(3, 3)
  %36 = f32[3,3]{1,0} aten::normal(%35, %33, %5), location=<module>@dropout_test.py:12
  %37 = f32[3,3]{1,0} aten::mul(%36, %31), [email protected]:958
  %38 = f32[3,3]{1,0} aten::mul(%37, %26), [email protected]:958
  %39 = f32[3,3]{1,0} aten::mul(%38, %21), [email protected]:958, ROOT=0
}

tensor([[-0.0000,  3.4327,  0.0000],
        [ 0.0000,  0.0000,  0.0000],
        [-1.6992, -0.0000, -0.0000]])

There is one S64 device data (the seed) and two-per-dropout F32 ones.
One of those is the probabilities of the dropout, but those should be cached (in the device data cache) after the first step ... unless they keep changing.

[davidel]

Changing the test to this:

import torch
import torch.nn as nn
import torch_xla
import torch_xla.core.xla_model as xm
import torch_xla.debug.metrics as met

device = xm.xla_device()

d1 = nn.Dropout(0.2)
d2 = nn.Dropout(0.25)
d3 = nn.Dropout(0.33)

for _ in range(0, 10):
  t = torch.randn(3, 3, device=device)
  x = d3(d2(d1(t)))

  print(torch_xla._XLAC._get_xla_tensors_text([x]))
  print(x.cpu())
  xm.mark_step()


print(met.metrics_report())

We can see that TransferToServerTime has 13 instances in a 10-long loop. Which shows that the seed is sent once per step, and the dropout device data scalars are cached:

Metric: TransferToServerTime
  TotalSamples: 13
  Accumulator: 040ms497.962us
  ValueRate: 043ms059.199us / second
  Rate: 13.8222 / second
  Percentiles: 1%=680.743us; 5%=680.743us; 10%=752.831us; 20%=753.868us; 50%=867.445us; 80%=001ms134.478us; 90%=008ms302.636us; 95%=023ms922.756us; 99%=023ms922.756us

[davidel]

Another option I was thinking is that the model's dropout probabilities are a function of the step number, which will make them change at every step.
But there has to be many dropout layers to make a dent in a non trivial step time.

[taylanbil]

My guess is, the code under scrutiny manually reset the seed

This indeed was the case. Eliminating this, the code is still about 1.5x slower compared to no dropout at all.

I'm going to be digging into this a bit more, comparing metrics in the "no dropout" vs "dropout+no set seed every step" cases. Will report here.

I will also check the dropout prob's to see if anything is dynamic there.

[taylanbil]

I verified that there is no change in dropout probability from step to step.

Here's a summary comparison b/w the dropout run and no dropout run. This skips setting seed manually every step.

                                  METRIC        NODROPOUT      DROPOUT       PCT CHANGE

           TransferToServerTime.Total_sec       28.5300      225.4089       690.1
                                 xla::mul         62715        140955       124.8
                 DeviceLockWait.Total_sec       47.2706        8.4059       -82.2
                    ExecuteTime.Total_sec      133.6597      230.9260        72.8
                     XrtReadLiteral.Count          2872           880       -69.4
         IrValueTensorToXlaData.Total_sec       16.3543        5.0458       -69.1
           XrtReleaseAllocation.Total_sec      954.0963      305.9116       -67.9
               TransferToServerTime.Count           739           980        32.6
                   TensorsGraphSize.Total  5939632.0000  6477952.0000         9.1
                        CompileTime.Count            15            15         0.0
                        ExecuteTime.Count           253           253         0.0
                                 MarkStep           246           246         0.0

lmk if I should upload the full metrics report.

Btw there seems to be 109 dropout layers in this workload.

[davidel]

So, from a pure execute POV, there is a slowdown (133.6597s vs 230.9260s) even w/out accounting TransferToServerTime.
This is a slowdown we noticed time ago WRT random number generation.
It might be worth feeding xprof traces of the two to the XLA team.
IIRC chatting with Blake there was a reshape which was creating problems. He did a CL to fix that time ago, but only helped some, and this might be something else.

Now to TransferToServerTime.
This is pretty bad. We seem to sending one extra (the seed) tensor per step, looking at the numbers (980 - 739 = 241 =~ 246 ... the step count).
But extra 241 scalar sends costed about 200s (225.4089 - 28.5300), which is ridiculous at about 1s per send.
There must be something going on. Is this repeatable in other VMs combos?

I started looking in making the scalar->device_data materializations asynchronous, which would help, but not if we are spending 1s to send a few bytes.
We can't hide 1s behind a few hundred ms step time 😉

As temp workaround which might be entertaining to try is to have a Dropout class which uses a single, shared, large, R1 F32 noise tensor (in [0, 1) domain), with the dropouts selecting a random window within it, reshape to tensor size, and do the compare.
But if XLA does not fuse the reshape it might not help.

[taylanbil]

Here's a small repro of this effect:

import torch
import torch.nn as nn
import torch_xla
import torch_xla.core.xla_model as xm
import torch_xla.debug.metrics as met

import numpy as np

device = xm.xla_device()

def metsumm(step=None):
    rep = met.metrics_report().split('\n')
    out = []
    prefix = '\tSTEP {}: '.format(step) if step else ''
    for i, line in enumerate(rep):
        if 'ExecuteTime' in line or 'TransferToServerTime' in line:
            out.extend([line,rep[i+1], rep[i+2]])
    for i in range(3):
        x = '{}{}'.format(prefix, out[i])
        x = x + ' '*(45-len(x)) + out[i+3].split(':')[-1]
        print(x)
    print('')


layers, N = [], 100
for _ in range(100):
    layers.extend([
        nn.Linear(N, N),
        nn.Dropout(np.random.uniform(0, 1)),
    ])
layers.append(nn.Linear(N, 1))
net = nn.Sequential(*layers).to(device)
t = torch.randn(N, N, device=device)
for i in range(0, 501):
  x = net(t)
  x.sum().backward()
  xm.mark_step()
  #print(torch_xla._XLAC._get_xla_tensors_text([x]))
  #print(x.cpu())
  if i and not (i % 50):
      metsumm(i)

        STEP 50: Metric: ExecuteTime                 TransferToServerTime
        STEP 50:   TotalSamples: 50                  353
        STEP 50:   Accumulator: 386ms819.130us       471ms997.083us

        STEP 100: Metric: ExecuteTime                TransferToServerTime
        STEP 100:   TotalSamples: 100                403
        STEP 100:   Accumulator: 695ms799.342us      572ms337.407us

        STEP 150: Metric: ExecuteTime                TransferToServerTime
        STEP 150:   TotalSamples: 150                453
        STEP 150:   Accumulator: 980ms343.245us      681ms668.675us

        STEP 200: Metric: ExecuteTime                TransferToServerTime
        STEP 200:   TotalSamples: 200                503
        STEP 200:   Accumulator: 01s274ms332.079us   786ms756.882us

        STEP 250: Metric: ExecuteTime                TransferToServerTime
        STEP 250:   TotalSamples: 250                553
        STEP 250:   Accumulator: 02s570ms918.918us   889ms583.393us

        STEP 300: Metric: ExecuteTime                TransferToServerTime
        STEP 300:   TotalSamples: 300                603
        STEP 300:   Accumulator: 02s855ms584.276us   01s003ms345.956us

        STEP 350: Metric: ExecuteTime                TransferToServerTime
        STEP 350:   TotalSamples: 350                653
        STEP 350:   Accumulator: 02s133ms378.060us   01s107ms554.303us

        STEP 400: Metric: ExecuteTime                TransferToServerTime
        STEP 400:   TotalSamples: 400                703
        STEP 400:   Accumulator: 02s408ms093.992us   01s224ms838.175us

        STEP 450: Metric: ExecuteTime                TransferToServerTime
        STEP 450:   TotalSamples: 450                753
        STEP 450:   Accumulator: 03s689ms412.681us   01s340ms450.411us

        STEP 500: Metric: ExecuteTime                TransferToServerTime
        STEP 500:   TotalSamples: 500                803
        STEP 500:   Accumulator: 03s965ms858.948us   01s456ms103.197us

[davidel]

Yes, we can see there is one scalar->device_data send per step, but that's roughly 2ms each.
The test model above have a few ms step time itself, so in there, it looks pretty bad.
What are roBERTa and Transformer typical step time, at batch sizes you are testing with?

BTW: ExecuteTime Accumulator goes down from step 300 to step 350, which looks odd.

[taylanbil]

There are 2 types of steps in the model workload;

• fwd pass: ~550 ms
• fwd+bwd pass: ~1400 ms

[davidel]

@taylanbil I feel we are missing something.
In this post https://github.com/pytorch/xla/issues/2636#issuecomment-731391360 it seems we are taking about 1s for one TransferToServerTime (which is outrageous) while in your previous example we take less than 2ms (which is still a bit high, but in the ballpark).

I assume the numbers you mention in https://github.com/pytorch/xla/issues/2636#issuecomment-731391360 are relative to a real model, right?
Let's work with that, as it seems the one triggering the bad behavior.

Can you post the full metrics at the end of the two runs (w/ and w/out dropout)?
One idea is that the 1s we are seeing as average is due to a few long latency samples ...

[taylanbil]

I understand, there's something that's going on in the model workload that doesn't occur in the small repro, which is causing high TransferToServerTime, when enabling dropout.

I got a lot of metrics reports for you. Restarting the tpu, I captured metrics reports, every 10 steps, w/ and w/o dropout.
I also dumped the graph sitting on the model's output tensor, following your idea from the small repro to count scalars on CPU, but there seems to be nothing there?

graph.txt.zip
met_reports.tar.zip

[taylanbil]

The only difference b/w these runs is, returning x in this line.

Wrote a quick script to put data from the previous comment executetime and transfertoservertime side by side:

EXECUTETIME PER STEP, dropout first, nodropout last
STEP 0 35s705ms831.199us 32s741ms067.409us
STEP 1 01m01s337ms338.026us 54s834ms594.183us
STEP 2 01m28s815ms985.614us 01m16s905ms247.117us
STEP 3 02m54s275ms454.577us 02m38s994ms565.257us
STEP 4 03m33s140ms072.057us 02m07s867ms208.952us
STEP 5 03m59s690ms041.520us 02m28s138ms650.507us
STEP 6 03m25s204ms389.389us 03m50s278ms462.110us
STEP 7 04m52s803ms140.492us 03m12s369ms421.685us
STEP 8 04m18s418ms312.617us 04m35s509ms202.457us
STEP 9 05m46s028ms321.924us 04m57s478ms760.666us
STEP 10 05m12s515ms167.626us 04m19s699ms480.659us
STEP 11 06m38s000ms214.624us 05m41s744ms153.164us
STEP 12 06m05s593ms504.593us 05m03s878ms467.318us
STEP 13 07m31s095ms086.413us 05m25s006ms637.882us
STEP 14 07m59s643ms593.141us 06m48s986ms045.273us
STEP 15 07m24s016ms024.551us 06m09s253ms377.687us
STEP 16 08m51s627ms438.116us 07m31s340ms710.938us
STEP 17 08m17s285ms820.092us 07m53s426ms458.522us
STEP 18 09m44s890ms155.917us 07m16s599ms411.565us
STEP 19 09m11s388ms838.549us 08m39s645ms662.096us
STEP 20 10m37s743ms432.002us 08m60s889ms521.214us
STEP 21 10m03s272ms063.267us 08m22s032ms826.426us
STEP 22 10m30s870ms679.449us 09m44s125ms283.821us
STEP 23 11m56s363ms279.530us 09m06s361ms812.456us
STEP 24 11m24s974ms823.240us 09m29s358ms335.403us
STEP 25 12m49s360ms747.737us 10m51s610ms993.292us
STEP 26 12m16s788ms539.300us 10m13s691ms030.728us
STEP 27 13m42s266ms062.832us 11m35s814ms246.528us
STEP 28 13m09s719ms552.243us 11m57s958ms183.750us
STEP 29 14m36s128ms809.403us 11m20s954ms413.184us
STEP 30 14m02s533ms348.828us 12m41s279ms732.860us
STEP 31 14m28s194ms238.605us 12m03s389ms828.065us
STEP 32 15m55s651ms292.278us 12m26s517ms475.675us
STEP 33 15m21s127ms842.122us 13m48s640ms163.817us
STEP 34 16m49s637ms807.035us 13m11s564ms681.957us
STEP 35 16m14s136ms151.632us 14m32s806ms600.800us
STEP 36 17m41s586ms621.933us 14m54s929ms796.052us
STEP 37 17m07s334ms019.625us 14m16s053ms896.940us
STEP 38 18m34s751ms004.062us 15m38s185ms119.940us
STEP 39 18m01s239ms323.326us 15m01s184ms893.022us

TRANSFERTOSERVERTIME PER STEP, dropout first, nodropout last
STEP 0 40s554ms714.225us 11s204ms563.365us
STEP 10 01m07s554ms753.069us 15s303ms476.302us
STEP 20 02m32s347ms066.907us 19s665ms997.525us
STEP 30 02m59s250ms079.632us 23s795ms765.455us
STEP 40 02m25s045ms677.250us 26s151ms234.497us
STEP 50 03m51s986ms202.117us 30s283ms280.635us
STEP 60 03m17s903ms859.609us 34s630ms924.512us
STEP 70 04m44s796ms862.595us 38s845ms706.795us
STEP 80 04m10s748ms829.784us 41s249ms493.593us
STEP 90 05m37s570ms578.737us 45s420ms186.596us
STEP 100 05m01s441ms288.760us 49s663ms811.652us
STEP 110 05m28s290ms341.323us 52s266ms704.501us
STEP 120 06m54s166ms135.137us 56s650ms246.037us
STEP 130 06m21s853ms276.830us 60s808ms354.771us
STEP 140 07m47s748ms703.297us 01m03s176ms407.889us
STEP 150 07m13s532ms410.813us 01m07s276ms542.287us
STEP 160 08m37s094ms243.967us 01m11s631ms814.643us
STEP 170 08m03s300ms925.012us 01m15s764ms615.206us
STEP 180 08m29s234ms415.605us 01m18s103ms787.492us
STEP 190 09m56s887ms929.204us 01m22s209ms031.706us
STEP 200 09m21s755ms860.063us 01m26s615ms521.636us
STEP 210 10m48s617ms325.379us 01m30s752ms877.682us
STEP 220 10m14s504ms173.933us 02m33s003ms616.651us
STEP 230 11m40s270ms415.717us 02m37s121ms122.864us
STEP 240 11m06s818ms194.311us 02m41s543ms725.141us
STEP 250 12m32s578ms969.412us 02m45s529ms863.981us
STEP 260 12m57s355ms799.479us 02m48s751ms007.288us
STEP 270 12m24s094ms310.727us 02m52s880ms225.667us
STEP 280 13m50s895ms563.421us 02m55s158ms179.370us
STEP 290 13m17s519ms920.958us 02m59s176ms354.848us
STEP 300 14m41s330ms329.317us 02m03s566ms713.466us
STEP 310 14m08s188ms132.033us 02m07s689ms564.816us
STEP 320 15m34s056ms879.550us 02m10s060ms605.733us
STEP 330 15m01s725ms008.781us 02m14s108ms657.451us
STEP 340 15m27s550ms194.519us 02m17s349ms563.633us
STEP 350 16m52s291ms094.508us 02m21s312ms473.992us
STEP 360 16m18s985ms950.950us 02m25s588ms311.190us
STEP 370 17m45s995ms970.389us 02m29s714ms953.058us
STEP 380 17m11s554ms546.316us 03m32s980ms436.153us
STEP 390 18m37s100ms510.736us 03m36s048ms164.842us

def getval(line):
    return line.split(':')[-1].strip()


def extract(fn):
    with open(fn, 'r') as f:
        rep=f.readlines()
    ex, t = [], []
    for i, line in enumerate(rep):
        if 'ExecuteTime' in line:
            ex.append(getval(rep[i+2]))
        if 'TransferToServerTime' in line:
            t.append(getval(rep[i+2]))
    return ex, t

dexs, dtrs, ndexs, ndtrs = [], [], [], []
for i in range(10, 401, 10):
    dfn='dropout/step{}.txt'.format(i)
    ndfn='nodropout/step{}.txt'.format(i)
    dex, dtr =extract(dfn)
    ndex, ndtr=extract(ndfn)
    dexs.append(dex)
    dtrs.append(dtr)
    ndexs.append(ndex)
    ndtrs.append(ndtr)
print('EXECUTETIME PER STEP, dropout first, nodropout last')
for i, (dex, ndex) in enumerate(zip(dexs, ndexs)):
    prefix='STEP {}'.format(i)
    for l1,l2 in zip(dex, ndex):
        print(prefix, l1, l2)
print('')
print('TRANSFERTOSERVERTIME PER STEP, dropout first, nodropout last')
for i, (dtr, ndtr) in enumerate(zip(dtrs, ndtrs)):
    prefix='STEP {}'.format(i*10)
    for l1,l2 in zip(dtr, ndtr):
        print(prefix, l1, l2)

The humanization bug from before is visible here too btw..

[davidel]

One thing I see at step 300:

Metric: OutboundData
  TotalSamples: 2485
  Accumulator: 1.48GB
  ValueRate: 30.59KB / second
  Rate: 1.49026 / second
  Percentiles: 1%=4.00B; 5%=4.00B; 10%=4.00B; 20%=8.00B; 50%=8.00B; 80%=8.00B; 90%=24.00B; 95%=24.00B; 99%=512.50KB

Metric: TransferToServerTime
  TotalSamples: 2485
  Accumulator: 13m17s519ms920.958us
  ValueRate: 617ms753.926us / second
  Rate: 1.49026 / second
  Percentiles: 1%=001ms307.247us; 5%=002ms561.117us; 10%=002ms672.751us; 20%=002ms859.661us; 50%=475ms180.383us; 80%=532ms540.442us; 90%=939ms518.915us; 95%=945ms224.238us; 99%=965ms036.705us

There is one OutboundData posting for each TransferToServerTime.
We can see that the majority of the outbound data size is 8 bytes (looks like the seed), but with a 1:1 mapping I'd expect the percentiles for TransferToServerTime to track the ones of OutboundData.
Instead we can see that already the 50% jumps up to 475ms.

This seems to agree with the idea that a few of those "cheap" seed sends end up taking hundreds of milliseconds.

[taylanbil]

I think most sends are taking hundreds of miliseconds; Here are STEP 0 percentiles:

TransferToServerTime PER STEP, dropout first, nodropout last
STEP 0 |  COUNTS:  666 vs 625 | VALUES:  40s554ms714.225us vs 11s204ms563.365us
          1%=790.867us            1%=927.685us
          5%=893.679us            5%=001ms073.821us
          10%=001ms005.993us      10%=001ms176.252us
          20%=001ms160.685us      20%=001ms293.840us
          50%=002ms642.572us      50%=002ms888.765us
          80%=009ms395.321us      80%=014ms740.750us
          90%=028ms098.210us      90%=035ms608.386us
          95%=475ms200.129us      95%=047ms056.125us
          99%=953ms563.897us      99%=209ms480.121us

And here are step 110 and 120 percentiles, when the TransferToServerTime count has roughly doubled from 666

STEP 110 |  COUNTS:  1356 vs 875 | VALUES:  05m28s290ms341.323us vs 52s266ms704.501us
          1%=844.015us            1%=943.986us
          5%=001ms096.672us       5%=001ms113.133us
          10%=001ms324.093us      10%=001ms225.792us
          20%=002ms612.980us      20%=001ms404.995us
          50%=030ms490.586us      50%=002ms016.566us
          80%=500ms094.125us      80%=043ms483.544us
          90%=937ms752.718us      90%=193ms350.174us
          95%=947ms677.849us      95%=202ms369.544us
          99%=972ms875.109us      99%=749ms582.520us

STEP 120 |  COUNTS:  1418 vs 897 | VALUES:  06m54s166ms135.137us vs 56s650ms246.037us
         1%=908.355us    1%=943.986us
          5%=001ms223.729us       5%=001ms114.905us
          10%=001ms436.936us      10%=001ms228.806us
          20%=002ms683.333us      20%=001ms410.485us
          50%=470ms608.817us      50%=002ms022.383us
          80%=508ms362.246us      80%=047ms056.125us
          90%=939ms599.041us      90%=194ms924.862us
          95%=947ms779.001us      95%=203ms430.441us
          99%=972ms875.109us      99%=749ms582.520us

It can be seen that 400ms+ sends eat into the lower percentiles and get to 50 percentile somewhere b/w 110 and 120's step, where the counts double from the initial value. it seems like once the model gets going, almost all the transfertoservertime's are > 400ms.

Looking at OutboundData and TransferToServerTime over all the steps reveals that in the dropout case, we send data to device 3x times per step (after removing the initial counts)

every 10 steps (model steps, which is 5 MarkSteps in xla terms due to grad accumulation):

• 62 sends when enabling dropout vs 22 sends when no dropout, per 10 steps.
• 26 seconds of TransferToServerTime with dropout, vs 4seconds and change w/o dropout, per 10 steps.

Looking at the profiles, these sends are blocking execution too, so nothing is hidden behind anything it seems like.

OutboundData PER STEP, dropout first, nodropout last
STEP 0 |  COUNTS:  666 vs 625 | VALUES:  1.45GB vs 1.45GB
STEP 10 |  COUNTS:  729 vs 648 | VALUES:  1.45GB vs 1.45GB
STEP 20 |  COUNTS:  791 vs 670 | VALUES:  1.45GB vs 1.45GB
STEP 30 |  COUNTS:  854 vs 693 | VALUES:  1.45GB vs 1.45GB
STEP 40 |  COUNTS:  917 vs 716 | VALUES:  1.45GB vs 1.45GB
STEP 50 |  COUNTS:  980 vs 739 | VALUES:  1.45GB vs 1.45GB
STEP 60 |  COUNTS:  1042 vs 761 | VALUES:  1.45GB vs 1.45GB
STEP 70 |  COUNTS:  1105 vs 784 | VALUES:  1.46GB vs 1.46GB
STEP 80 |  COUNTS:  1167 vs 806 | VALUES:  1.46GB vs 1.46GB
STEP 90 |  COUNTS:  1231 vs 830 | VALUES:  1.46GB vs 1.46GB
STEP 100 |  COUNTS:  1293 vs 852 | VALUES:  1.46GB vs 1.46GB
STEP 110 |  COUNTS:  1356 vs 875 | VALUES:  1.46GB vs 1.46GB
STEP 120 |  COUNTS:  1418 vs 897 | VALUES:  1.46GB vs 1.46GB
STEP 130 |  COUNTS:  1481 vs 920 | VALUES:  1.46GB vs 1.46GB
STEP 140 |  COUNTS:  1544 vs 943 | VALUES:  1.46GB vs 1.46GB
STEP 150 |  COUNTS:  1607 vs 966 | VALUES:  1.47GB vs 1.47GB
STEP 160 |  COUNTS:  1669 vs 988 | VALUES:  1.47GB vs 1.47GB
STEP 170 |  COUNTS:  1732 vs 1011 | VALUES:  1.47GB vs 1.47GB
STEP 180 |  COUNTS:  1794 vs 1033 | VALUES:  1.47GB vs 1.47GB
STEP 190 |  COUNTS:  1858 vs 1057 | VALUES:  1.47GB vs 1.47GB
STEP 200 |  COUNTS:  1920 vs 1079 | VALUES:  1.47GB vs 1.47GB
STEP 210 |  COUNTS:  1983 vs 1102 | VALUES:  1.47GB vs 1.47GB
STEP 220 |  COUNTS:  2045 vs 1124 | VALUES:  1.47GB vs 1.47GB
STEP 230 |  COUNTS:  2108 vs 1147 | VALUES:  1.47GB vs 1.47GB
STEP 240 |  COUNTS:  2171 vs 1170 | VALUES:  1.48GB vs 1.48GB
STEP 250 |  COUNTS:  2234 vs 1193 | VALUES:  1.48GB vs 1.48GB
STEP 260 |  COUNTS:  2296 vs 1215 | VALUES:  1.48GB vs 1.48GB
STEP 270 |  COUNTS:  2359 vs 1238 | VALUES:  1.48GB vs 1.48GB
STEP 280 |  COUNTS:  2421 vs 1260 | VALUES:  1.48GB vs 1.48GB
STEP 290 |  COUNTS:  2485 vs 1284 | VALUES:  1.48GB vs 1.48GB
STEP 300 |  COUNTS:  2547 vs 1306 | VALUES:  1.48GB vs 1.48GB
STEP 310 |  COUNTS:  2610 vs 1329 | VALUES:  1.48GB vs 1.48GB
STEP 320 |  COUNTS:  2672 vs 1351 | VALUES:  1.49GB vs 1.49GB
STEP 330 |  COUNTS:  2735 vs 1374 | VALUES:  1.49GB vs 1.49GB
STEP 340 |  COUNTS:  2798 vs 1397 | VALUES:  1.49GB vs 1.49GB
STEP 350 |  COUNTS:  2861 vs 1420 | VALUES:  1.49GB vs 1.49GB
STEP 360 |  COUNTS:  2923 vs 1442 | VALUES:  1.49GB vs 1.49GB
STEP 370 |  COUNTS:  2986 vs 1465 | VALUES:  1.49GB vs 1.49GB
STEP 380 |  COUNTS:  3048 vs 1487 | VALUES:  1.49GB vs 1.49GB
STEP 390 |  COUNTS:  3112 vs 1511 | VALUES:  1.49GB vs 1.49GB
----------------------------------------
TransferToServerTime PER STEP, dropout first, nodropout last
STEP 0 |  COUNTS:  666 vs 625 | VALUES:  40s554ms714.225us vs 11s204ms563.365us
STEP 10 |  COUNTS:  729 vs 648 | VALUES:  01m07s554ms753.069us vs 15s303ms476.302us
STEP 20 |  COUNTS:  791 vs 670 | VALUES:  02m32s347ms066.907us vs 19s665ms997.525us
STEP 30 |  COUNTS:  854 vs 693 | VALUES:  02m59s250ms079.632us vs 23s795ms765.455us
STEP 40 |  COUNTS:  917 vs 716 | VALUES:  02m25s045ms677.250us vs 26s151ms234.497us
STEP 50 |  COUNTS:  980 vs 739 | VALUES:  03m51s986ms202.117us vs 30s283ms280.635us
STEP 60 |  COUNTS:  1042 vs 761 | VALUES:  03m17s903ms859.609us vs 34s630ms924.512us
STEP 70 |  COUNTS:  1105 vs 784 | VALUES:  04m44s796ms862.595us vs 38s845ms706.795us
STEP 80 |  COUNTS:  1167 vs 806 | VALUES:  04m10s748ms829.784us vs 41s249ms493.593us
STEP 90 |  COUNTS:  1231 vs 830 | VALUES:  05m37s570ms578.737us vs 45s420ms186.596us
STEP 100 |  COUNTS:  1293 vs 852 | VALUES:  05m01s441ms288.760us vs 49s663ms811.652us
STEP 110 |  COUNTS:  1356 vs 875 | VALUES:  05m28s290ms341.323us vs 52s266ms704.501us
STEP 120 |  COUNTS:  1418 vs 897 | VALUES:  06m54s166ms135.137us vs 56s650ms246.037us
STEP 130 |  COUNTS:  1481 vs 920 | VALUES:  06m21s853ms276.830us vs 60s808ms354.771us
STEP 140 |  COUNTS:  1544 vs 943 | VALUES:  07m47s748ms703.297us vs 01m03s176ms407.889us
STEP 150 |  COUNTS:  1607 vs 966 | VALUES:  07m13s532ms410.813us vs 01m07s276ms542.287us
STEP 160 |  COUNTS:  1669 vs 988 | VALUES:  08m37s094ms243.967us vs 01m11s631ms814.643us
STEP 170 |  COUNTS:  1732 vs 1011 | VALUES:  08m03s300ms925.012us vs 01m15s764ms615.206us
STEP 180 |  COUNTS:  1794 vs 1033 | VALUES:  08m29s234ms415.605us vs 01m18s103ms787.492us
STEP 190 |  COUNTS:  1858 vs 1057 | VALUES:  09m56s887ms929.204us vs 01m22s209ms031.706us
STEP 200 |  COUNTS:  1920 vs 1079 | VALUES:  09m21s755ms860.063us vs 01m26s615ms521.636us
STEP 210 |  COUNTS:  1983 vs 1102 | VALUES:  10m48s617ms325.379us vs 01m30s752ms877.682us
STEP 220 |  COUNTS:  2045 vs 1124 | VALUES:  10m14s504ms173.933us vs 02m33s003ms616.651us
STEP 230 |  COUNTS:  2108 vs 1147 | VALUES:  11m40s270ms415.717us vs 02m37s121ms122.864us
STEP 240 |  COUNTS:  2171 vs 1170 | VALUES:  11m06s818ms194.311us vs 02m41s543ms725.141us
STEP 250 |  COUNTS:  2234 vs 1193 | VALUES:  12m32s578ms969.412us vs 02m45s529ms863.981us
STEP 260 |  COUNTS:  2296 vs 1215 | VALUES:  12m57s355ms799.479us vs 02m48s751ms007.288us
STEP 270 |  COUNTS:  2359 vs 1238 | VALUES:  12m24s094ms310.727us vs 02m52s880ms225.667us
STEP 280 |  COUNTS:  2421 vs 1260 | VALUES:  13m50s895ms563.421us vs 02m55s158ms179.370us
STEP 290 |  COUNTS:  2485 vs 1284 | VALUES:  13m17s519ms920.958us vs 02m59s176ms354.848us
STEP 300 |  COUNTS:  2547 vs 1306 | VALUES:  14m41s330ms329.317us vs 02m03s566ms713.466us
STEP 310 |  COUNTS:  2610 vs 1329 | VALUES:  14m08s188ms132.033us vs 02m07s689ms564.816us
STEP 320 |  COUNTS:  2672 vs 1351 | VALUES:  15m34s056ms879.550us vs 02m10s060ms605.733us
STEP 330 |  COUNTS:  2735 vs 1374 | VALUES:  15m01s725ms008.781us vs 02m14s108ms657.451us
STEP 340 |  COUNTS:  2798 vs 1397 | VALUES:  15m27s550ms194.519us vs 02m17s349ms563.633us
STEP 350 |  COUNTS:  2861 vs 1420 | VALUES:  16m52s291ms094.508us vs 02m21s312ms473.992us
STEP 360 |  COUNTS:  2923 vs 1442 | VALUES:  16m18s985ms950.950us vs 02m25s588ms311.190us
STEP 370 |  COUNTS:  2986 vs 1465 | VALUES:  17m45s995ms970.389us vs 02m29s714ms953.058us
STEP 380 |  COUNTS:  3048 vs 1487 | VALUES:  17m11s554ms546.316us vs 03m32s980ms436.153us
STEP 390 |  COUNTS:  3112 vs 1511 | VALUES:  18m37s100ms510.736us vs 03m36s048ms164.842us

[cjolivier01]

Just FYI, if you send a true constant as the seed down to RngUniform (random.cpp), it won't be used, but you can do something like this to use the old-style stateful RNG that won't require an xfer at all:

xla::XlaOp RngUniform(xla::XlaOp seed, const xla::Shape& shape,
                      xla::XlaOp minval, xla::XlaOp maxval) {
  if (seed.builder()->IsConstant(seed).ValueOrDie()) {
    return xla::RngUniform(minval, maxval, shape);
  }
  xla::XlaOp rng_seed = MakeSeed(seed);
  xla::Shape rng_shape = MakeRngShape(shape);
  xla::XlaOp rng_minval = MakeUniformBoundaryValue(minval);
...

(in order to baseline if there's a difference)

[davidel]

The seed is not constant, otherwise dropouts will drop the same elements at every step.

[davidel]

@taylanbil 400+ms is huge, but in theory there should be only one extra device data being sent WRT no-dropout.
That extra device data should be the seed, as the eventual other recurring constants should be cached (unless we are hitting the cache limit - XLA_DEVDATA_CACHE_SIZE to bump it up from 128).
Have you tested this on a different GVM/RVM sets (maybe in another cell)?

[cjolivier01]

The seed is not constant, otherwise dropouts will drop the same elements at every step.

right , I just meant if you force a constant in there or simply call the other one, as a test maybe can isolate if the delay is indeed that particular send of that seed, since it won't be sent in that case.

[taylanbil]

ok, let me try w/ a bigger XLA_DEVDATA_CACHE_SIZE.

This was tested on several setups (versions, VM size, etc) but perhaps not on a different zone. I'll try a run on us-east1-c or something and report.

[taylanbil]

XLA_DEVDATA_CACHE_SIZE=512 helped a tiny bit (~1%), not much. compared to no dropout, still huge difference.

[taylanbil]

@davidel Today, @ultrons has reproed the delta w/ and w/o dropout, in us-central1-c. This seems to reliably happen everywhere so far.

In order to partition this problem into two, I ran a version of this workload with just one dropout layer. The comparison of metrics are below:

• no dropout vs one dropout layer
• one dropout layer vs full (109) dropout layers

Quick takeaways:

• Performance is significantly degraded w/ one dropout, but not as bad as full dropout:
  
  
  
  • 33% slow w/ one dropout layer, 61% slow w/ full (109 layers) (both compared to no dropout)
  
  
• For some weird reason, TransferToServerTime counts are exactly the same b/w one and full dropout, but it takes less time w/ one dropout.
• ExecuteTime becomes 20% slower w/ full dropout compared to one dropout layer. Could this be a separate problem than the rng seed, perhaps layout etc related?

Here are all metrics reports collected so far: met_reports-new.tar.zip

[davidel]

Thanks @taylanbil!
Am I right the model under test is doing one optimizer step for two mark-steps?
The STEP in your debug metrics are the optimizer ones or the mark-step ones?
It looks like there are 4 device data sends for each one of your STEP AFAICT, between no-dropout and one-dropout.
This is odd. Does the dropout probability change for each step?
How does the one-dropout code look like?

It might help enable IR debug and correlate the IR graph to see where the extra device-data are attached to.

[davidel]

The extra ExecuteTime is the potential XLA perf "issue" which I mentioned somewhere in this thread, related to the performance of the XLA bit-generator.

[taylanbil]

this does 3 fwd passes and 1 fwd+bwd pass (optimizer step is this one), per model's step.
we do mark step after the non-final fwd passes and after the fwd+bwd.

So, the STEP is the model's step, and MarkStep values corresponding to a given STEP are about 4x (we also mark step before logging, so that's an additional markstep per log).

I'm fairly certain we're not changing the dropout probability. In order to completely verify, I added the following diff (which does one dropout only, with fixed proba):

--- a/fairseq/modules/fairseq_dropout.py
+++ b/fairseq/modules/fairseq_dropout.py
@@ -11,6 +11,7 @@ import torch.nn.functional as F


 logger = logging.getLogger(__name__)
+FIRST=True


 class FairseqDropout(nn.Module):
@@ -20,9 +21,16 @@ class FairseqDropout(nn.Module):
         self.p = p
         self.module_name = module_name
         self.apply_during_inference = False
+        global FIRST
+        self._first = FIRST
+        FIRST = False

     def forward(self, x, inplace: bool = False):
-        if self.training or self.apply_during_inference:
+        if self._first and self.training or self.apply_during_inference:
+            import torch_xla.core.xla_model as xm
+            print('DROPPING OUT', xm.get_ordinal())
+            return F.dropout(x, p=0.1, training=True, inplace=inplace)
             return F.dropout(x, p=self.p, training=True, inplace=inplace)
         else:
             return x

[taylanbil]

this way, it was a tad slower, I made the following modification and it's the same speed as doing one dropout:

--- a/fairseq/modules/fairseq_dropout.py
+++ b/fairseq/modules/fairseq_dropout.py
@@ -11,6 +11,7 @@ import torch.nn.functional as F


 logger = logging.getLogger(__name__)
+FIRST=True


 class FairseqDropout(nn.Module):
@@ -18,11 +19,20 @@ class FairseqDropout(nn.Module):
     def __init__(self, p, module_name=None):
         super().__init__()
         self.p = p
+        self.pppp = 0.1
         self.module_name = module_name
         self.apply_during_inference = False
+        global FIRST
+        self._first = FIRST
+        FIRST = False

     def forward(self, x, inplace: bool = False):
-        if self.training or self.apply_during_inference:
+        #if self.training or self.apply_during_inference:
+        if self._first and self.training or self.apply_during_inference:
+            import torch_xla.core.xla_model as xm
+            print('DROPPING OUT', xm.get_ordinal())
+            return F.dropout(x, p=self.pppp, training=True, inplace=inplace)
+            return F.dropout(x, p=0.1, training=True, inplace=inplace)
             return F.dropout(x, p=self.p, training=True, inplace=inplace)
         else:
             return x

[davidel]

Just to be clear, during one STEP we are ingesting 4 samples, right?
In any case the 4x is explained then. What makes the seed to reset is the MarkStep.

So the real question is why the transfer to such small data ends up taking so long (400+ms).
Here we show a 5x time for a 2x count:

STEP 390 |  COUNTS:  3112 vs 1511 | VALUES:  18m37s100ms510.736us vs 03m36s048ms164.842us

The core server side metric for device data allocation reports sane values, so this is very likely network:

Metric: XrtAllocateFromTensor
  TotalSamples: 63262
  Accumulator: 02m48s760ms362.480us
  Mean: 002ms827.791us
  StdDev: 944.729us
  Rate: 47.1327 / second
  Percentiles: 25%=598.797us; 50%=002ms127.860us; 80%=003ms641.138us; 90%=003ms828.182us; 95%=003ms956.698us; 99%=003ms357.947us

[taylanbil]

In any case the 4x is explained then

yes, it is 1 per Markstep, sorry for the confusion there.

Yeah XrtAllocateFromTensor does not increase much. Ok so let's say it's network, how would we approach a solution? Would background bulk seed sends (instead of 1 per step, move seed forward on client 1000 times and send all, and churn through the seeds on the client as we mark step) work?

[taylanbil]

@davidel The XLA team identified a sub-optimality, and we're seeing some gains in performance after patching the TPU runtime w/ a fix. I'll be collecting and posting new xla metrics soon. Not sure if this will give us a completely different picture or not, but it may be worth looking at that set of metrics before discussing further...

[taylanbil]

The patched runtime metrics are attached here. Summary here. As expected, executetime improved. Perhaps unexpected, transfertoservertime also improved, but the delta remains:

ExecuteTime PER STEP, ORDER: ['nodropout', 'onedropout', 'dropout']
STEP 340 |  COUNTS:  1498 vs 1498 vs 1498 | VALUES:  10m24s866ms902.847us vs 11m38s911ms573.812us vs 12m20s414ms710.975us
TransferToServerTime PER STEP, ORDER: ['nodropout', 'onedropout', 'dropout']
STEP 340 |  COUNTS:  1397 vs 2798 vs 2798 | VALUES:  02m32s698ms015.149us vs 10m33s091ms288.556us vs 11m22s155ms713.784us

Prior to the patch, TransferToServerTime used to be

ExecuteTime PER STEP, ORDER: ['nodropout', 'onedropout', 'dropout']
STEP 390 |  COUNTS:  1511 vs 3112 vs 3112 | VALUES:  03m36s048ms164.842us vs 14m21s370ms066.914us vs 18m37s100ms510.736us

So we're doing 314 less sends and it's saving 4mins for one dropout case and 7 mins for full.

Could it be that the TransferToServerTime metric is actually mis-reporting it and hiding what's the real difference?

met_reports-patched.tar.zip

[taylanbil]

btw, other notable differences b/w no-dropout and one-dropout are:

python /usr/share/torch-xla-nightly/pytorch/xla/scripts/metrics_compare.py nodropout/step50.txt onedropout/step50.txt -t 0 --no-humanize --topn-counters 1000 --topn-percentiles 1000 | grep -v Xrt
...
                                      KEY          Val1          Val2  PCT_CHANGE

...
             ReleaseDataHandlesTime.Count          2384          4217        76.9
                 DeviceLockWait.Total_sec       37.8460       14.9282       -60.6
         IrValueTensorToXlaData.Total_sec       12.9352        5.3660       -58.5
...

[taylanbil]

The previous 2 comments were w/ forking the rng seed every time. I forgot that piece of code in there. Removing that, the metrics reports look similar, so those comments aren't way off, but below are the new metrics w/o forking the rng seed every time, so it removes one more complication from the measurements:

summary gist
met_reports-patched-nofork.tar.zip

ExecuteTime PER STEP, ORDER: ['nodropout', 'onedropout', 'dropout']
330 |  COUNTS:  1443 vs 1443 vs 1443 | VALUES:  10m57s375ms591.032us vs 10m05s120ms506.367us vs 12m44s940ms259.590us
TransferToServerTime PER STEP, ORDER: ['nodropout', 'onedropout', 'dropout']
STEP 340 |  COUNTS:  1397 vs 2798 vs 2798 | VALUES:  02m40s553ms119.882us vs 09m29s306ms942.357us vs 11m18s321ms907.027us

Prior to the patch, TransferToServerTime used to be

ExecuteTime PER STEP, ORDER: ['nodropout', 'onedropout', 'dropout']
STEP 390 |  COUNTS:  1511 vs 3112 vs 3112 | VALUES:  03m36s048ms164.842us vs 14m21s370ms066.914us vs 18m37s100ms510.736us

Other notable differences, no dropout first, one-dropout 2nd:

                                      KEY          Val1          Val2  PCT_CHANGE
...
                 DeviceLockWait.Total_sec       32.5173        2.0075       -93.8
             ReleaseDataHandlesTime.Count          2470          4022        62.8
         IrValueTensorToXlaData.Total_sec       14.6748        5.5019       -62.5
  TransferToServerTransformTime.Total_sec        0.8906        0.5513       -38.1

[davidel]

The notable differences are in line with the behavior.
The 2nd, 3rd and 4th are due to creating more data handles (seed -> device_data).
The 1st is due to the model overall step time being slower, and hence we wait less at the barrier.

[taylanbil]

so, is the current thinking that this is network latency? how can we progress towards a solution to this?

[davidel]

The TransferToServerTime is the total time seen by the client.
We know that the core XrtAllocateFromTensor is sane, so it must be a SW/HW layer (lower level client side, network, lower level service side) underneath.
Given the amount of them (400+ms) I am more inclined to point to network (and for this I mean anything underneath the RVM/GVM guests interconnects).

[davidel]

Maybe create a small micro-bench with a few small tensors and a tiny computation executed many step (with new tensor values), and check the TransferToServerTime?

[taylanbil]

Not sure if this is exactly what you want, but here's my microbench:

import torch
import torch.nn as nn
import torch_xla
import torch_xla.core.xla_model as xm
import torch_xla.distributed.parallel_loader as pl
import torch_xla.debug.metrics as met

import numpy as np

def metsumm(step=None):
    rep = met.metrics_report().split('\n')
    out = []
    prefix = '\tSTEP {}: '.format(step) if step else ''
    for i, line in enumerate(rep):
        if 'ExecuteTime' in line or 'TransferToServerTime' in line:
            out.extend([line,rep[i+1], rep[i+2]])
    for i in range(3):
        x = '{}{}'.format(prefix, out[i])
        x = x + ' '*(45-len(x)) + out[i+3].split(':')[-1]
        print(x)
    print('\n')


class loadergen(torch.utils.data.Dataset):

    def __init__(self, N, nmodes):
        self.n = N
        self.nmodes = nmodes

    def __len__(self):
        return 400

    def __getitem__(self, j):
        return [torch.randn(self.n, self.n) for _ in range(self.nmodes)]


class Net(nn.Module):

    def __init__(self, nmodes, inpsize, dropout):
        super().__init__()
        self.nmodes = nmodes
        self.inpsize = inpsize
        self.dropout = dropout
        self.setup()

    def setup(self):
        self.modes = []
        for _ in range(self.nmodes):
            mode = [nn.Linear(self.inpsize, 1)]
            if self.dropout:
                mode.append(nn.Dropout(0.2))
            self.modes.append(nn.Sequential(*mode))
        self.modes = nn.ModuleList(self.modes)

    def forward(self, inp):
        return sum(
            layer(i) for layer, i in zip(self.modes, inp)
        )

def main(drp):
    INPSIZE, NMODES = 2, 100
    loader = loadergen(INPSIZE, NMODES)
    device = xm.xla_device()
    loader = pl.MpDeviceLoader(loader, device)
    net = Net(NMODES, INPSIZE, drp).to(device)
    for i, t in enumerate(loader):
        if i > 400:
            break
        x = net(t)
        x.sum().backward()
        xm.mark_step()
        if i and not (i % 50):
            metsumm(i)

import sys
main(sys.argv[1]=='dropout')

Except for the humanization bug, I'm not seeing a wildly unexpected item here. WDYT?

        STEP 50: Metric: ExecuteTime                 TransferToServerTime
        STEP 50:   TotalSamples: 51                  266
        STEP 50:   Accumulator: 462ms879.572us       02s892ms054.781us


        STEP 100: Metric: ExecuteTime                TransferToServerTime
        STEP 100:   TotalSamples: 101                329
        STEP 100:   Accumulator: 855ms375.998us      03s770ms303.921us


        STEP 150: Metric: ExecuteTime                TransferToServerTime
        STEP 150:   TotalSamples: 151                391
        STEP 150:   Accumulator: 01s236ms357.997us   04s522ms036.929us


        STEP 200: Metric: ExecuteTime                TransferToServerTime
        STEP 200:   TotalSamples: 201                454
        STEP 200:   Accumulator: 02s587ms274.229us   04s283ms226.547us


        STEP 250: Metric: ExecuteTime                TransferToServerTime
        STEP 250:   TotalSamples: 251                516
        STEP 250:   Accumulator: 02s942ms570.149us   05s991ms282.915us


        STEP 300: Metric: ExecuteTime                TransferToServerTime
        STEP 300:   TotalSamples: 301                579
        STEP 300:   Accumulator: 02s281ms558.972us   06s729ms705.564us


        STEP 350: Metric: ExecuteTime                TransferToServerTime
        STEP 350:   TotalSamples: 351                641
        STEP 350:   Accumulator: 03s622ms366.526us   06s429ms008.171us


        STEP 400: Metric: ExecuteTime                TransferToServerTime
        STEP 400:   TotalSamples: 401                704
        STEP 400:   Accumulator: 03s974ms729.015us   07s183ms415.301us

--------------

NO DROPOUT

        STEP 50: Metric: ExecuteTime                 TransferToServerTime
        STEP 50:   TotalSamples: 51                  214
        STEP 50:   Accumulator: 562ms649.154us       01s250ms315.113us


        STEP 100: Metric: ExecuteTime                TransferToServerTime
        STEP 100:   TotalSamples: 101                227
        STEP 100:   Accumulator: 01s048ms016.370us   02s897ms408.677us


        STEP 150: Metric: ExecuteTime                TransferToServerTime
        STEP 150:   TotalSamples: 151                239
        STEP 150:   Accumulator: 02s517ms252.211us   02s489ms330.526us


        STEP 200: Metric: ExecuteTime                TransferToServerTime
        STEP 200:   TotalSamples: 201                252
        STEP 200:   Accumulator: 02s989ms697.050us   03s116ms576.429us


        STEP 250: Metric: ExecuteTime                TransferToServerTime
        STEP 250:   TotalSamples: 251                264
        STEP 250:   Accumulator: 02s441ms725.324us   04s693ms187.980us


        STEP 300: Metric: ExecuteTime                TransferToServerTime
        STEP 300:   TotalSamples: 301                277
        STEP 300:   Accumulator: 03s915ms750.514us   04s334ms137.957us


        STEP 350: Metric: ExecuteTime                TransferToServerTime
        STEP 350:   TotalSamples: 351                289
        STEP 350:   Accumulator: 03s378ms809.434us   05s922ms025.266us


        STEP 400: Metric: ExecuteTime                TransferToServerTime
        STEP 400:   TotalSamples: 401                302
        STEP 400:   Accumulator: 04s849ms769.362us   06s558ms596.463us

[davidel]

Hmm, it seems here the TransferToServerTime is under control. 0.7s for 400 extra seed sends ... sounds about right.

What are the specs of the RVM driving the real test? Any chance is under overload?

[taylanbil]

I'm using an e2-standard-32 lately, I tried w/ n1-highmem-96 as well, and @ultrons tried w/ n1-standard-32.

Right now, I'm running on e2-st-32, and here's top:

top - 19:46:42 up 35 days, 20:20,  5 users,  load average: 2.86, 2.15, 2.71
Tasks: 369 total,   1 running, 368 sleeping,   0 stopped,   0 zombie
%Cpu(s): 10.1 us,  0.1 sy,  0.0 ni, 89.8 id,  0.0 wa,  0.0 hi,  0.1 si,  0.0 st
KiB Mem : 13203813+total, 84613392 free, 45924592 used,  1500152 buff/cache
KiB Swap:        0 total,        0 free,        0 used. 84783088 avail Mem

  PID USER      PR  NI    VIRT    RES    SHR S  %CPU %MEM     TIME+ COMMAND
25145 taylanb+  20   0 14.466g 5.448g 235508 S  43.0  4.3   6:00.84 python
25146 taylanb+  20   0 14.097g 5.401g 235096 S  40.1  4.3   6:04.78 python
25150 taylanb+  20   0 14.167g 5.395g 235756 S  40.1  4.3   5:57.14 python
25151 taylanb+  20   0 14.238g 5.392g 236068 S  42.1  4.3   5:55.99 python
25147 taylanb+  20   0 14.171g 5.391g 235472 S  39.7  4.3   6:06.17 python
25149 taylanb+  20   0 14.167g 5.391g 235660 S  42.1  4.3   6:05.99 python
25152 taylanb+  20   0 14.234g 5.391g 234672 S  40.4  4.3   5:55.13 python
25148 taylanb+  20   0 14.167g 5.390g 234844 S  41.1  4.3   3:25.06 python
26037 taylanb+  20   0 2315068 388452 220424 S   0.0  0.3   0:03.07 python
26052 taylanb+  20   0 2315048 387896 219880 S   0.0  0.3   0:03.12 python
26038 taylanb+  20   0 2315060 387824 219800 S   0.0  0.3   0:03.11 python
26039 taylanb+  20   0 2315064 387748 219728 S   0.0  0.3   0:03.08 python
26042 taylanb+  20   0 2315072 387700 219664 S   0.0  0.3   0:03.06 python
26050 taylanb+  20   0 2315068 387472 219444 S   0.0  0.3   0:03.05 python
26046 taylanb+  20   0 2315004 387308 219292 S   0.0  0.3   0:03.06 python
26040 taylanb+  20   0 2315064 387276 219244 S   0.0  0.3   0:03.05 python
25139 taylanb+  20   0 2169644 385304 217292 S   0.0  0.3   0:04.24 python

and here's free-h

$ free -h
              total        used        free      shared  buff/cache   available
Mem:           125G         43G         80G        229M        1.4G         80G
Swap:            0B          0B          0B

[davidel]

I find it a bit interesting that the 400ms or so "long times" we see in the TransferToServerTime are close to the ExecuteTime (in the real test).
That could in theory be explained by an interaction between TransferToServer and Execute.
We know that does not look like a service side core (TPU runtime) locking, as the XRT metric looks good.
So (assuming the interaction exists) it could be anywhere from GRPC down.

I am not sure whether enabling GRPC tracing on the client side can reveal useful info:

https://github.com/grpc/grpc/blob/master/doc/environment_variables.md

BTW, forgot to ask. Are the metrics above single core or 8 cores (can we get single core ones if this is the case)?

[taylanbil]

apologies for not replying for a while, I will be trying the suggestions next week. The metrics I posted are from an 8 core run (4 way modelparallel, 2 way dataparallel)

I observed the effect on a single core run as well, although I would have to collect metrics.

[taylanbil]

I still have not gotten to trying the suggestions from your last comment @davidel, but here's a result from another attempt:

here's a diff by @JackCaoG that sends XLA_SEED_CACHE_SIZE seeds to device every time, and rng runs thru those across many steps.

Trying w/ this, I see tremendous improvement (30% on top of the improvement seen in https://github.com/pytorch/xla/issues/2636#issuecomment-738988155) in e2e wall time for the workload in hand.

[JackCaoG]

That diff is a quick hack, I think if we want to make this change to the master, we need to

1. Handle the manual seed better. With this change every time a RNG_SEED is set by user, it will send XLA_SEED_CACHE_SIZE new seeds to the server, which is unnecessary.
2. Invalidate the used seed

[davidel]

As I mentioned (I think 😄 ) before, that is probably the nicest fix we can do.
We can make device data materialization async, but we do have a problem somewhere (the 400+ms sends) which we need to look at, as this is likely some problem with our system (from client up to service infrastructure).
I feel solving that will be solving the original model, as scalar sends should be around 1ms, not 400+.