Distributed TPU Training, training data stored in GCS

[tottenjordan]

We have built a terraform script that spins up 4 VMs and uses a v3-32 TPU for Resnet50 training. We store the Imagenet training and validation data in a GCS bucket. Full code repo can be found here

• we use the torch_xla.distributed.xla_dist
• as well as the test_train_mp_imagenet.py (only altering it to use our GCS data loader)

For the questions below, I've attached a log file (with metrics_debug), and used the following configuration:

• VM machine types = n2-custom (72 vCPUs, 512 GB memory)
• NUM_EPOCHS=20
• BATCH_SIZE=512
• TEST_BATCH_SIZE=64
• NUM_WORKERS=8
• log_steps=200
• --conda-env=torch-xla-1.7
• --env XLA_USE_BF16=1
• • default learning rate and schedule

Questions

• Not sure what baseline to compare with, but epoch training time seems to be around 5-6 minutes.
  • This is true for 8 workers on batch sizes of 128, 256, and 512. (batch size of 128 with 32 workers seems to be low 4 minutes per epoch).
  • Is there anything from a code or configuration perspective we could do to improve this? 32 workers seems like overkill, but we've seen better results with this?
• Sometime we will get BrokenPipeError: [Errno 32] Broken pipe or unhealthy mesh errors and training will automatically restart (see line 20689 in log file for Broken Pipe Error during Epoch 13).
  • Is there anything we can do to overcome this?

imagenetraw_logfiles4-v3-32-512batch-8workers.txt

@zcain117
@shanemhansen

[tottenjordan]

Here is a screenshot of the dashboard for the previously mentioned training job

[taylanbil]

how does the xla metrics report look? Does it differ from our regular nightly resnet50 runs too much in terms of ExecuteTime, CompileTime, etc ? @zcain117

[zcain117]

• Compiles looks fine - no new compiles after the 1st epoch
• no aten:: calls except for reasonable amount of aten::_local_scalar_dense

I don't have any recent logs with metrics so I will wrap up a few of my experiments running now and then try to kick off 2 runs on v3-32, one using regular PD and one with SSD. Then we can compare metrics

[zcain117]

Attached the metrics for my v3-32 run (filtered to just 1 of the 4 VM's stdout)
512_ssd_logs_grepped.txt

• CompileTime is static for both after epoch 1
• ExecuteTime yours: 04m56s -> 05m26s -> 07m35s -> 08m15s -> 09m45s -> 11m03s -> 12m23s -> 14m52s -> 16m13s -> 18m38s -> 19m49s -> 21m48s -> 22m01s -> 23m20s (epoch 14)
• ExecuteTime mine: 02m06s -> 03m13s -> 04m18s -> 05m16s -> 06m21s -> 07m30s -> 09m34s -> 10m35s -> 11m34s -> 12m40s -> 13m45s -> 14m44s -> 15m15s (epoch 14)
• No major difference in TransferToServerTime or TransferFromServerTime
• DeviceLockWait total was ~11m on mine and ~20m on yours

Some of the XrtMetrics:

• Metric: XrtAllocateFromTensor.c_tpu_worker.0 was ~1hr01m on mine and ~1hr26m on yours
• The XrtExecute was substantially higher on your run. Here is from my run after 14 epochs:

2020-12-17 20:20:12 10.164.0.108 [0] Metric: XrtExecute.c_tpu_worker.0
2020-12-17 20:20:12 10.164.0.108 [0]   TotalSamples: 65744
2020-12-17 20:20:12 10.164.0.108 [0]   Accumulator: 05h10m41s479ms268.855us
2020-12-17 20:20:12 10.164.0.108 [0]   Mean: 490ms928.087us
2020-12-17 20:20:12 10.164.0.108 [0]   StdDev: 475ms630.074us
2020-12-17 20:20:12 10.164.0.108 [0]   Rate: 8.95582 / second
2020-12-17 20:20:12 10.164.0.108 [0]   Percentiles: 25%=030ms152.881us; 50%=532ms559.099us; 80%=839ms410.966us; 90%=941ms093.544us; 95%=01s107ms439.003us; 99%=02s541ms791.543us
2020-12-17 20:20:12 10.164.0.108 [0] Metric: XrtExecute.c_tpu_worker.1
2020-12-17 20:20:12 10.164.0.108 [0]   TotalSamples: 65744
2020-12-17 20:20:12 10.164.0.108 [0]   Accumulator: 05h07m12s957ms501.585us
2020-12-17 20:20:12 10.164.0.108 [0]   Mean: 497ms052.484us
2020-12-17 20:20:12 10.164.0.108 [0]   StdDev: 489ms446.525us
2020-12-17 20:20:12 10.164.0.108 [0]   Rate: 8.97735 / second
2020-12-17 20:20:12 10.164.0.108 [0]   Percentiles: 25%=030ms993.646us; 50%=533ms292.432us; 80%=834ms846.378us; 90%=01s008ms699.441us; 95%=01s223ms020.681us; 99%=02s710ms407.093us
2020-12-17 20:20:12 10.164.0.108 [0] Metric: XrtExecute.c_tpu_worker.2
2020-12-17 20:20:12 10.164.0.108 [0]   TotalSamples: 65744
2020-12-17 20:20:12 10.164.0.108 [0]   Accumulator: 05h02m44s926ms805.700us
2020-12-17 20:20:12 10.164.0.108 [0]   Mean: 478ms717.034us
2020-12-17 20:20:12 10.164.0.108 [0]   StdDev: 446ms347.398us
2020-12-17 20:20:12 10.164.0.108 [0]   Rate: 8.97536 / second
2020-12-17 20:20:12 10.164.0.108 [0]   Percentiles: 25%=030ms261.579us; 50%=550ms147.834us; 80%=803ms044.583us; 90%=886ms119.476us; 95%=01s028ms439.400us; 99%=02s965ms099.527us
2020-12-17 20:20:12 10.164.0.108 [0] Metric: XrtExecute.c_tpu_worker.3
2020-12-17 20:20:12 10.164.0.108 [0]   TotalSamples: 65744
2020-12-17 20:20:12 10.164.0.108 [0]   Accumulator: 05h60m37s069ms776.382us
2020-12-17 20:20:12 10.164.0.108 [0]   Mean: 492ms107.899us
2020-12-17 20:20:12 10.164.0.108 [0]   StdDev: 466ms978.928us
2020-12-17 20:20:12 10.164.0.108 [0]   Rate: 8.95665 / second
2020-12-17 20:20:12 10.164.0.108 [0]   Percentiles: 25%=030ms009.936us; 50%=572ms417.431us; 80%=779ms379.190us; 90%=893ms028.297us; 95%=01s192ms527.232us; 99%=02s026ms606.519us

And yours:

19309 2020-12-16 16:35:11 10.164.0.15 [0] Metric: XrtExecute.c_tpu_worker.0
19310 2020-12-16 16:35:11 10.164.0.15 [0]   TotalSamples: 36793
19311 2020-12-16 16:35:11 10.164.0.15 [0]   Accumulator: 13h49m38s272ms501.708us
19312 2020-12-16 16:35:11 10.164.0.15 [0]   Mean: 555ms735.382us
19313 2020-12-16 16:35:11 10.164.0.15 [0]   StdDev: 958ms486.022us
19314 2020-12-16 16:35:11 10.164.0.15 [0]   Rate: 2.88988 / second
19315 2020-12-16 16:35:11 10.164.0.15 [0]   Percentiles: 25%=029ms421.710us; 50%=496ms488.162us; 80%=558ms749.285us; 90%=693ms959.466us; 95%=02s444ms      354.965us; 99%=06s568ms853.819us
19316 2020-12-16 16:35:11 10.164.0.15 [0] Metric: XrtExecute.c_tpu_worker.1
19317 2020-12-16 16:35:11 10.164.0.15 [0]   TotalSamples: 36794
19318 2020-12-16 16:35:11 10.164.0.15 [0]   Accumulator: 13h31m54s695ms601.280us
19319 2020-12-16 16:35:11 10.164.0.15 [0]   Mean: 555ms451.089us
19320 2020-12-16 16:35:11 10.164.0.15 [0]   StdDev: 956ms374.188us
19321 2020-12-16 16:35:11 10.164.0.15 [0]   Rate: 2.88073 / second
19322 2020-12-16 16:35:11 10.164.0.15 [0]   Percentiles: 25%=030ms713.184us; 50%=496ms333.137us; 80%=555ms457.000us; 90%=687ms633.420us; 95%=03s646ms      049.735us; 99%=05s333ms692.363us
19323 2020-12-16 16:35:11 10.164.0.15 [0] Metric: XrtExecute.c_tpu_worker.2
19324 2020-12-16 16:35:11 10.164.0.15 [0]   TotalSamples: 36794
19325 2020-12-16 16:35:11 10.164.0.15 [0]   Accumulator: 13h40m56s431ms993.757us
19326 2020-12-16 16:35:11 10.164.0.15 [0]   Mean: 553ms892.010us
19327 2020-12-16 16:35:11 10.164.0.15 [0]   StdDev: 885ms607.958us
19328 2020-12-16 16:35:11 10.164.0.15 [0]   Rate: 2.88393 / second
19329 2020-12-16 16:35:11 10.164.0.15 [0]   Percentiles: 25%=030ms643.081us; 50%=497ms562.074us; 80%=561ms749.341us; 90%=679ms752.559us; 95%=03s819ms      256.936us; 99%=05s577ms093.633us
19330 2020-12-16 16:35:11 10.164.0.15 [0] Metric: XrtExecute.c_tpu_worker.3
19331 2020-12-16 16:35:11 10.164.0.15 [0]   TotalSamples: 36793
19332 2020-12-16 16:35:11 10.164.0.15 [0]   Accumulator: 13h45m05s942ms659.191us
19333 2020-12-16 16:35:11 10.164.0.15 [0]   Mean: 582ms085.747us
19334 2020-12-16 16:35:11 10.164.0.15 [0]   StdDev: 01s044ms268.753us
19335 2020-12-16 16:35:11 10.164.0.15 [0]   Rate: 2.87976 / second
19336 2020-12-16 16:35:11 10.164.0.15 [0]   Percentiles: 25%=030ms526.811us; 50%=497ms784.653us; 80%=559ms869.838us; 90%=682ms877.519us; 95%=03s790ms      202.928us; 99%=06s726ms888.469us

@JackCaoG do you think XrtExecute is usually an important metric?

@tottenjordan have you compared your speed if using a PD or SSD PD instead of the custom GCS reader? And regarding the custom GCS reader, if there any way you could increase the prefetching amount?

[taylanbil]

XrtExecute doesn't get reset b/w runs, you need to start tpu again for it to be reset. This looks like it didn't get restarted b/w runs?

ExecuteTime is very different. Are the models/batch sizes etc exactly the same?

[zcain117]

GCS was ~5m30s per epoch and SSD PD was ~1m30s per epoch so ~4min * 14 epochs = ~56min cumulative difference but ExecuteTime was only ~8min cumulative difference. Probably something else is causing the majority of the difference?

[zcain117]

Forgot to include my training command:
python -m torch_xla.distributed.xla_dist --tpu=$SSD_TPU_NAME --conda-env=torch-xla-1.7 --env XLA_USE_BF16=1 --env ANY_OTHER=ENV_VAR -- python /usr/share/torch-xla-1.7/pytorch/xla/test/test_train_mp_imagenet.py --model=resnet50 --num_workers=8 --batch_size=512 --log_steps=200 --num_epochs=14 --metrics_debug --datadir=/mnt/disks/dataset/imagenet 2>&1 | tee ssd_training_logs.txt

[zcain117]

@tottenjordan I think the next thing to try is increase the prefetch_factor when you initialize the DataLoader here. Maybe try a few different sizes (default value is 2)

I think this change should be in the 1.7 version you're using: https://github.com/pytorch/pytorch/releases/tag/v1.7.0

[taylanbil]

what are the percentiles for executetime?

[tottenjordan]

@zcain117 , @shanemhansen ran training jobs with PD and will post log/metrics report when available.

Here is first successful run with prefetch_factor=4. Attached log/metrics report, but here are some highlights:

• v2-32 (due to TPU capacity constraints)
• reduced batch size from 512 to 256
  • (v2-32 couldnt handle 512 + prefetch of 4*8 workers)
• average epoch training time is 6:24
  • Doesn't include epoch 1 of 6:38
  • epoch 5 training time = 9:21; epoch 10 = 10:24

export NUM_EPOCHS=14
export BATCH_SIZE=256
export TEST_BATCH_SIZE=64
export NUM_WORKERS=8

python -m torch_xla.distributed.xla_dist --tpu=$TPU_NAME --conda-env=torch-xla-1.7 \
    --env XLA_USE_BF16=1 \
    -- python /tmp/thepackage/test_train_mp_imagenet.py \
    --num_epochs=$NUM_EPOCHS \
    --batch_size=$BATCH_SIZE \
    --num_workers=$NUM_WORKERS \
    --log_steps=200 \
    --logdir=$LOGDIR \
    --datadir=$IMAGE_DIR \
    --test_set_batch_size=$TEST_BATCH_SIZE \
    --metrics_debug

imagenetr1aw_logfiles6-v2-32-256batch-8wrks-4prefetch.txt

@zcain117 @taylanbil Our data are stored as JPEG. Is it worth exploring tfrecords or other formats? Are the data in your PD tests stored as JPEG?

Will run prefetch_factor=8 next

[tottenjordan]

prefetch_factor=8

• v2-32
• average epoch training time = 5:43
  • doesnt include epoch 1 of 6:07
  • more consistent and lower average than prefetch_factor=4 w/ 8 workers

export NUM_EPOCHS=14
export BATCH_SIZE=256
export TPU_NAME=mytpu
export TEST_BATCH_SIZE=64
export NUM_WORKERS=8

python -m torch_xla.distributed.xla_dist --tpu=$TPU_NAME --conda-env=torch-xla-1.7 \
    --env XLA_USE_BF16=1 \
    -- python /tmp/thepackage/test_train_mp_imagenet.py \
    --num_epochs=$NUM_EPOCHS \
    --batch_size=$BATCH_SIZE \
    --num_workers=$NUM_WORKERS \
    --log_steps=200 \
    --logdir=$LOGDIR \
    --datadir=$IMAGE_DIR \
    --test_set_batch_size=$TEST_BATCH_SIZE \
    --metrics_debug

imagenetr1aw_logfiles7-v2-32-256batch-8wrks-8prefetch.txt

[tottenjordan]

This is interesting. Increased to num_workers=16 and kept prefetch_factor=8

• v2-32
• average epoch training is 3:53
  • consistent throughout job

export NUM_EPOCHS=10
export BATCH_SIZE=256
export TEST_BATCH_SIZE=64
export NUM_WORKERS=16

python -m torch_xla.distributed.xla_dist --tpu=$TPU_NAME --conda-env=torch-xla-1.7 \
    --env XLA_USE_BF16=1 \
    -- python /tmp/thepackage/test_train_mp_imagenet.py \
    --num_epochs=$NUM_EPOCHS \
    --batch_size=$BATCH_SIZE \
    --num_workers=$NUM_WORKERS \
    --log_steps=200 \
    --logdir=$LOGDIR \
    --datadir=$IMAGE_DIR \
    --test_set_batch_size=$TEST_BATCH_SIZE \
    --metrics_debug

imagenetr1aw_logfiles8-v2-32-256batch-16wrks-8prefetch.txt

@zcain117 how do i get metrics_debug to print out in sequence by worker, instead of every other line potentially being from another worker?

[zcain117]

I don't know of any way to order the outputs. I just filter to 1 of the VM's output at a time when reading.

E.g. you see lines like this where the VM's IP is mentioned: 2020-12-18 15:52:56 10.128.0.70 [2] Step 0

So I was using grep "10.128.0.70" logs.txt

Earlier you mentioned good results with num_workers=32, maybe it's worth trying that + prefetch_factor even though that many workers is higher than we normally recommend

[tottenjordan]

Switched to v3-32 TPU pod.

Increasing the num_workers and prefetch_factor is showing better results:

Simply, moving from v2-32 to v3-32 showed a slight speed up of ~20s

• num_workers=16 and prefetch_factor=8
• average epoch training time = 3:32
• metrics_debug: imagenetraw_logfiles9-v3-32-256batch-16workers-8prefetch.txt

On v3-32, increasing num_workers and keeping prefetch_factor constant, showed speed-up of ~23s
num_workers=32 and prefetch_factor=8

• average epoch training time = 3:09
• metrics_debug: imagenetraw_logfiles10-v3-32-256batch-32workers-8prefetch.txt

2020-12-21 16:29:15 10.164.0.29 [0]   TotalSamples: 448
2020-12-21 16:29:15 10.164.0.29 [0]   Accumulator: 276ms908.139us
2020-12-21 16:29:15 10.164.0.29 [0]   Mean: 615.866us
2020-12-21 16:29:15 10.164.0.29 [0]   StdDev: 001ms449.249us
2020-12-21 16:29:15 10.164.0.29 [0]   Rate: 0.0584148 / second
2020-12-21 16:29:15 10.164.0.29 [0]   Percentiles: 25%=350.820us; 50%=398.093us; 80%=497.008us; 90%=616.153us; 95%=937.838us; 99%=007ms552.195us
2020-12-21 16:29:15 10.164.0.29 [0] Metric: XrtReadLiteral.c_tpu_worker.1
2020-12-21 16:29:15 10.164.0.29 [0]   TotalSamples: 448
2020-12-21 16:29:15 10.164.0.29 [0]   Accumulator: 225ms417.282us
2020-12-21 16:29:15 10.164.0.29 [0]   Mean: 503.164us
2020-12-21 16:29:15 10.164.0.29 [0]   StdDev: 842.045us
2020-12-21 16:29:15 10.164.0.29 [0]   Rate: 0.0584197 / second
2020-12-21 16:29:15 10.164.0.29 [0]   Percentiles: 25%=352.787us; 50%=406.847us; 80%=508.414us; 90%=598.991us; 95%=718.260us; 99%=002ms539.905us
2020-12-21 16:29:15 10.164.0.29 [0] Metric: XrtReadLiteral.c_tpu_worker.2
2020-12-21 16:29:15 10.164.0.29 [0]   TotalSamples: 448
2020-12-21 16:29:15 10.164.0.29 [0]   Accumulator: 261ms891.467us
2020-12-21 16:29:15 10.164.0.29 [0]   Mean: 582.347us
2020-12-21 16:29:15 10.164.0.29 [0]   StdDev: 001ms256.640us
2020-12-21 16:29:15 10.164.0.29 [0]   Rate: 0.0584187 / second
2020-12-21 16:29:15 10.164.0.29 [0]   Percentiles: 25%=352.927us; 50%=411.384us; 80%=496.829us; 90%=583.585us; 95%=760.846us; 99%=006ms572.937us
2020-12-21 16:29:15 10.164.0.29 [0] Metric: XrtReadLiteral.c_tpu_worker.3
2020-12-21 16:29:15 10.164.0.29 [0]   TotalSamples: 448
2020-12-21 16:29:15 10.164.0.29 [0]   Accumulator: 257ms218.651us
2020-12-21 16:29:15 10.164.0.29 [0]   Mean: 574.149us
2020-12-21 16:29:15 10.164.0.29 [0]   StdDev: 001ms380.042us
2020-12-21 16:29:15 10.164.0.29 [0]   Rate: 0.0584215 / second
2020-12-21 16:29:15 10.164.0.29 [0]   Percentiles: 25%=355.008us; 50%=405.369us; 80%=516.042us; 90%=597.709us; 95%=727.514us; 99%=004ms020.203us```