Loss backward took forever until memory leaking on TPU v3

username: bryan
cooked: <p>duplicated post in Github. Please go to <a href="https://github.com/pytorch/xla/issues/3203#issuecomment-963769873" rel="noopener nofollow ugc">here</a> for more details.</p>
<p>I am fine tuning XLNet for multi-label text classification on a TPU v3. When the labels are under 100, the training went well and took about 10s per iterator(batch size 16). But if the labels are about 9000, the training stuck after the first batch forward at loss.backward. I tried to keep it running for about 4 hours until memory leaking. The model has about 131032033 parameters. The following is the metric reports.</p>
<p>Metric: CompileTime<br>
TotalSamples: 21<br>
Accumulator: 594ms960.157us<br>
ValueRate: 092ms946.360us / second<br>
Rate: 3.25085 / second<br>
Percentiles: 1%=001ms048.724us; 5%=001ms143.656us; 10%=001ms297.037us; 20%=002ms503.725us; 50%=002ms686.961us; 80%=005ms211.937us; 90%=006ms457.997us; 95%=018ms385.346us; 99%=529ms270.667us<br>
Metric: DeviceLockWait<br>
TotalSamples: 21<br>
Accumulator: 096.075us<br>
ValueRate: 013.724us / second<br>
Rate: 2.99983 / second<br>
Percentiles: 1%=003.520us; 5%=003.626us; 10%=003.792us; 20%=003.934us; 50%=004.301us; 80%=005.171us; 90%=005.555us; 95%=005.833us; 99%=007.407us<br>
Metric: ExecuteTime<br>
TotalSamples: 21<br>
Accumulator: 185ms133.715us<br>
ValueRate: 029ms301.158us / second<br>
Rate: 3.32368 / second<br>
Percentiles: 1%=002ms989.937us; 5%=002ms014.393us; 10%=002ms107.005us; 20%=002ms137.575us; 50%=002ms313.633us; 80%=003ms652.595us; 90%=003ms785.078us; 95%=005ms872.769us; 99%=136ms111.333us<br>
Metric: InboundData<br>
TotalSamples: 20<br>
Accumulator: 160.00B<br>
ValueRate: 31.90B / second<br>
Rate: 3.98692 / second<br>
Percentiles: 1%=8.00B; 5%=8.00B; 10%=8.00B; 20%=8.00B; 50%=8.00B; 80%=8.00B; 90%=8.00B; 95%=8.00B; 99%=8.00B<br>
Metric: InputOutputAliasCount<br>
TotalSamples: 21<br>
Accumulator: 232.00<br>
ValueRate: 33.19 / second<br>
Rate: 3.00399 / second<br>
Percentiles: 1%=1.00; 5%=1.00; 10%=1.00; 20%=1.00; 50%=1.00; 80%=1.00; 90%=1.00; 95%=1.00; 99%=212.00<br>
Metric: IrValueTensorToXlaData<br>
TotalSamples: 232<br>
Accumulator: 05s153ms259.368us<br>
ValueRate: 566ms014.728us / second<br>
Rate: 25.482 / second<br>
Percentiles: 1%=830.401us; 5%=889.505us; 10%=001ms033.071us; 20%=001ms173.958us; 50%=004ms322.203us; 80%=015ms407.568us; 90%=036ms242.661us; 95%=085ms422.525us; 99%=355ms438.993us<br>
Metric: OutboundData<br>
TotalSamples: 239<br>
Accumulator: 984.20MB<br>
ValueRate: 61.39MB / second<br>
Rate: 14.9075 / second<br>
Percentiles: 1%=4.00B; 5%=3.00KB; 10%=3.00KB; 20%=3.00KB; 50%=12.00KB; 80%=2.25MB; 90%=9.00MB; 95%=14.74MB; 99%=48.00MB<br>
Metric: ReleaseDataHandlesTime<br>
TotalSamples: 21<br>
Accumulator: 027ms167.485us<br>
ValueRate: 004ms305.010us / second<br>
Rate: 3.3277 / second<br>
Percentiles: 1%=507.093us; 5%=512.067us; 10%=564.039us; 20%=604.354us; 50%=684.380us; 80%=925.953us; 90%=004ms156.686us; 95%=005ms773.718us; 99%=006ms807.402us<br>
Metric: TensorsGraphSize<br>
TotalSamples: 21<br>
Accumulator: 1420.00<br>
ValueRate: 219.98 / second<br>
Rate: 3.25324 / second<br>
Percentiles: 1%=18.00; 5%=18.00; 10%=18.00; 20%=18.00; 50%=18.00; 80%=18.00; 90%=18.00; 95%=18.00; 99%=1060.00<br>
Metric: TransferFromServerTime<br>
TotalSamples: 20<br>
Accumulator: 031ms137.437us<br>
ValueRate: 006ms207.132us / second<br>
Rate: 3.98693 / second<br>
Percentiles: 1%=864.648us; 5%=937.089us; 10%=948.000us; 20%=986.957us; 50%=001ms342.661us; 80%=002ms528.235us; 90%=003ms378.802us; 95%=004ms301.261us; 99%=004ms301.261us<br>
Metric: TransferToServerTime<br>
TotalSamples: 239<br>
Accumulator: 05s345ms883.669us<br>
ValueRate: 340ms073.866us / second<br>
Rate: 15.2066 / second<br>
Percentiles: 1%=820.348us; 5%=879.701us; 10%=001ms020.336us; 20%=001ms163.830us; 50%=004ms304.142us; 80%=017ms207.125us; 90%=041ms485.225us; 95%=090ms678.567us; 99%=355ms415.786us<br>
Metric: TransferToServerTransformTime<br>
TotalSamples: 239<br>
Accumulator: 769ms568.847us<br>
ValueRate: 048ms935.319us / second<br>
Rate: 14.9063 / second<br>
Percentiles: 1%=048.697us; 5%=054.600us; 10%=059.095us; 20%=079.733us; 50%=280.509us; 80%=002ms895.032us; 90%=003ms721.550us; 95%=015ms214.894us; 99%=052ms474.712us<br>
Counter: CreateCompileHandles<br>
Value: 21<br>
Counter: CreateDataHandles<br>
Value: 501<br>
Counter: CreateXlaTensor<br>
Value: 139157<br>
Counter: DestroyDataHandles<br>
Value: 232<br>
Counter: DestroyXlaTensor<br>
Value: 78849<br>
Counter: DeviceDataCacheMiss<br>
Value: 4<br>
Counter: MarkStep<br>
Value: 1<br>
Counter: ReleaseDataHandles<br>
Value: 232<br>
Counter: UncachedCompile<br>
Value: 21<br>
Counter: XRTAllocateFromTensor_Empty<br>
Value: 35<br>
Counter: XrtCompile_Empty<br>
Value: 32<br>
Counter: XrtExecuteChained_Empty<br>
Value: 32<br>
Counter: XrtExecute_Empty<br>
Value: 32<br>
Counter: XrtMemoryInfo_Empty<br>
Value: 32<br>
Counter: XrtRead_Empty<br>
Value: 32<br>
Counter: XrtReleaseAllocationHandle_Empty<br>
Value: 32<br>
Counter: XrtReleaseCompileHandle_Empty<br>
Value: 32<br>
Counter: XrtSessionCount<br>
Value: 4<br>
Counter: XrtSubTuple_Empty<br>
Value: 32<br>
Counter: aten::<em>local_scalar_dense<br>
Value: 20<br>
Counter: xla::<em>copy_from<br>
Value: 272<br>
Counter: xla::<em>softmax<br>
Value: 5161<br>
Counter: xla::<em>unsafe_view<br>
Value: 15363<br>
Counter: xla::add<br>
Value: 851<br>
Counter: xla::add</em><br>
Value: 243<br>
Counter: xla::addcmul<br>
Value: 240<br>
Counter: xla::arange_out<br>
Value: 120<br>
Counter: xla::as_strided<br>
Value: 222<br>
Counter: xla::bernoulli</em><br>
Value: 513<br>
Counter: xla::binary_cross_entropy_with_logits<br>
Value: 1<br>
Counter: xla::bmm<br>
Value: 6241<br>
Counter: xla::clamp_min</em><br>
Value: 1<br>
Counter: xla::div</em><br>
Value: 513<br>
Counter: xla::embedding<br>
Value: 10<br>
Counter: xla::empty<br>
Value: 916<br>
Counter: xla::empty_strided<br>
Value: 222<br>
Counter: xla::exp_<br>
Value: 2<br>
Counter: xla::expand<br>
Value: 10082<br>
Counter: xla::fill_<br>
Value: 1<br>
Counter: xla::gelu<br>
Value: 120<br>
Counter: xla::gt<br>
Value: 20<br>
Counter: xla::index_select<br>
Value: 130<br>
Counter: xla::log_<br>
Value: 1<br>
Counter: xla::max<br>
Value: 10<br>
Counter: xla::mean<br>
Value: 1<br>
Counter: xla::min<br>
Value: 10<br>
Counter: xla::mm<br>
Value: 10322<br>
Counter: xla::mul<br>
Value: 754<br>
Counter: xla::mul_<br>
Value: 1<br>
Counter: xla::native_batch_norm<br>
Value: 240<br>
Counter: xla::ne<br>
Value: 10<br>
Counter: xla::neg<br>
Value: 13<br>
Counter: xla::permute<br>
Value: 6130<br>
Counter: xla::rsub<br>
Value: 10<br>
Counter: xla::scatter_<br>
Value: 10<br>
Counter: xla::select<br>
Value: 5071<br>
Counter: xla::sigmoid<br>
Value: 1<br>
Counter: xla::slice<br>
Value: 10692<br>
Counter: xla::squeeze<br>
Value: 10062<br>
Counter: xla::stack<br>
Value: 5<br>
Counter: xla::sub<br>
Value: 121<br>
Counter: xla::sub_<br>
Value: 1<br>
Counter: xla::sum<br>
Value: 1<br>
Counter: xla::t<br>
Value: 10322<br>
Counter: xla::tanh<br>
Value: 5041<br>
Counter: xla::transpose<br>
Value: 5076<br>
Counter: xla::unsqueeze<br>
Value: 3670<br>
Counter: xla::view<br>
Value: 41789<br>
Counter: xla::zero_<br>
Value: 10<br>
Metric: XrtAllocateFromTensor<br>
TotalSamples: 2166<br>
Accumulator: 05s371ms945.919us<br>
Mean: 003ms247.374us<br>
StdDev: 005ms051.251us<br>
Rate: 65.5897 / second<br>
Percentiles: 25%=342.407us; 50%=002ms514.806us; 80%=005ms992.474us; 90%=007ms063.137us; 95%=015ms325.965us; 99%=019ms539.374us<br>
Metric: XrtCompile<br>
TotalSamples: 168<br>
Accumulator: 05s519ms861.698us<br>
Mean: 027ms897.986us<br>
StdDev: 106ms354.463us<br>
Rate: 5.10487 / second<br>
Percentiles: 25%=167.783us; 50%=197.449us; 80%=242.641us; 90%=159ms036.450us; 95%=162ms310.616us; 99%=169ms269.848us<br>
Metric: XrtExecute<br>
TotalSamples: 168<br>
Accumulator: 237ms422.502us<br>
Mean: 001ms413.229us<br>
StdDev: 478.713us<br>
Rate: 5.10719 / second<br>
Percentiles: 25%=001ms202.048us; 50%=001ms314.841us; 80%=001ms478.719us; 90%=002ms594.452us; 95%=002ms117.205us; 99%=004ms608.056us<br>
Metric: XrtExecutorEvict<br>
TotalSamples: 0<br>
Accumulator: nanB<br>
Mean: nanB<br>
StdDev: nanB<br>
Percentiles:<br>
Metric: XrtReadLiteral<br>
TotalSamples: 160<br>
Accumulator: 073ms435.633us<br>
Mean: 458.973us<br>
StdDev: 261.592us<br>
Rate: 4.89464 / second<br>
Percentiles: 25%=351.912us; 50%=432.780us; 80%=531.511us; 90%=579.884us; 95%=659.538us; 99%=728.484us<br>
Metric: XrtReleaseAllocation<br>
TotalSamples: 205<br>
Accumulator: 007ms322.887us<br>
Mean: 035.721us<br>
StdDev: 038.990us<br>
Rate: 6.23232 / second<br>
Percentiles: 25%=023.236us; 50%=027.079us; 80%=034.452us; 90%=046.374us; 95%=082.995us; 99%=267.722us</p>

username: ElsebaiyMohamed
cooked: <p>There is any sloution for this issue?</p>
<p>i face the same problem with pytorch lightning, i tried defferent  batch size but still not work.<br>
i can train the model with batch size 512 with max num of tokens = 350 on kaggle gpu P100 but stuck on backward step of first batch when try TPU</p>